% This must be in the first 5 lines to tell arXiv to use pdfLaTeX, which is strongly recommended.
\pdfoutput=1
% In particular, the hyperref package requires pdfLaTeX in order to break URLs across lines.

\documentclass[twocolumn]{article}
\usepackage{authblk}
% Remove the "review" option to generate the final version.
\usepackage[final]{ACL2023}

% Standard package includes
\usepackage{times}
\usepackage{latexsym}

% For proper rendering and hyphenation of words containing Latin characters (including in bib files)
\usepackage[T1]{fontenc}
% For Vietnamese characters
% \usepackage[T5]{fontenc}
% See https://www.latex-project.org/help/documentation/encguide.pdf for other character sets

% This assumes your files are encoded as UTF8
\usepackage[utf8]{inputenc}

% This is not strictly necessary, and may be commented out.
% However, it will improve the layout of the manuscript,
% and will typically save some space.
\usepackage{microtype}

% This is also not strictly necessary, and may be commented out.
% However, it will improve the aesthetics of text in
% the typewriter font.
\usepackage{inconsolata}

%=================================
\usepackage{graphicx}
\usepackage{caption}
\usepackage{subcaption}
\usepackage{adjustbox,booktabs}
\usepackage{makecell}
\usepackage{amsmath} 
\usepackage{dblfloatfix}
\usepackage{tabularx}
%\usepackage{subfigure}
\usepackage{multirow}
\usepackage{soul}
%\usepackage{xcolor}
\usepackage{booktabs}
% Standard package includes
\usepackage{times}
\usepackage{latexsym}
%\usepackage{authblk}
\usepackage[table]{xcolor} 
%\usepackage{authblk}

\renewcommand\theadalign{bc}
\renewcommand\theadfont{\bfseries}
\renewcommand\theadgape{\Gape[4pt]}
\renewcommand\cellgape{\Gape[4pt]}

\usepackage{tikz}
\def\checkmark{\tikz\fill[scale=0.4](0,.35) -- (.25,0) -- (1,.7) -- (.25,.15) -- cycle;} 
\newcommand{\tikzxmark}{%
\tikz[scale=0.23] {
    \draw[line width=0.7,line cap=round] (0,0) to [bend left=6] (1,1);
    \draw[line width=0.7,line cap=round] (0.2,0.95) to [bend right=3] (0.8,0.05);
}}
\newcommand{\centered}[1]{\begin{tabular}{l} #1 \end{tabular}}
%=================================
\newcommand{\ignore}[1]{}
\newcommand\pascale[1]{\textcolor{red}{#1}}
\newcommand\leila[1]{\textcolor{green}{#1}}
% If the title and author information does not fit in the area allocated, uncomment the following
%
% \setlength\titlebox{<dim>}
%
% and set <dim> to something 5cm or larger.
\title{Flatness-Aware Gradient Descent for Safe Conversational AI \vspace{1.2em}}

% \author{Leila Khalatbari \\
%   HKUST/Hong Kong \\ 
%   \texttt{lkhalatbari@connect.ust.hk} \\ \And
%   Saeid Hosseini \\
%   Sohar University/Oman \\
%   \texttt{SaHosseini@su.edu.om} \\ \And 
%     Hossein Sameti \\
%   SUT/Iran \\
%   \texttt{sameti@sharif.edu} \\ \And
%     Pascale FUNG \\
%   Meta\footnote{The work has been done when she was at HKUST}/France\\}

\author[1,3]{\bf \large Leila Khalatbari}
\author[2]{\bf \large Saeid Hosseini}
\author[3]{\bf \large Hossein Sameti}
\author[1]{\bf \large Pascale Fung}

\affil[1]{Hong Kong University of Science and Technology, Hong Kong}
\affil[2]{Sohar University, Oman}
\affil[3]{Sharif University of Technology, Iran \vspace{0.7em}}
% \affil[ ]{}
\affil[ ]{lkhalatbari@connect.ust.hk}
% \affil[ ]{\small SaHosseini@su.edu.om, \small sameti@sharif.edu}
% \affil[ ]{\small sameti@sharif.edu, \small pascale@ece.ust.hk}
% \affil[ ]{\small pascale@ece.ust.hk}

% \author[1]{Author A}
% \author[1]{Author B}
% \author[1]{Author C}
% \author[2]{Author D}
% \author[2]{Author E}
% \affil[1]{Department of Computer Science, \LaTeX\ University}
% \affil[2]{Department of Mechanical Engineering, \LaTeX\ University}
% \affil[ ]{\textit {\{email1,email2,email3,email4,email5\}@xyz.edu}}

% \author{%
%   Leila Khalatbari\thanks{\hspace{1mm} These authors contributed equally.} \\
%   HKUST\\
%   Hong Kong\\
%   \texttt{lkhalatbari@connect.ust.hk} \\
%   % examples of more authors
%   \And
%   Saeid Hosseini$^*$ \\
%   Sohar University \\
%   Oman \\
%   \texttt{SaHosseini@su.edu.om} \\
%   \AND
%   Hossein Sameti$^*$ \\
%   SUT \\
%   Iran \\
%   \texttt{sameti@sharif.edu} \\
%   \And
%   Pascale FUNG \\
%   Meta\footnote{The work has been done when she was at HKUST} \\
%   }

\begin{document}
{\makeatletter\acl@finalcopytrue
  \maketitle
}
\begin{abstract}
 As generative dialog models become ubiquitous in real-world applications, it is paramount to ensure a harmless generation. There are two major challenges when enforcing safety to open-domain chatbots. Firstly, it is impractical to provide training data reflecting the desired response to all emerging forms of toxicity (generalisation challenge). Secondly, implementing safety features may compromise the quality of the conversation (trade-off challenge). To tackle the challenges, this paper introduces a regularized fine-tuning approach called FlatGD. By employing a safety-tailored loss, we translate better optimization to more safety. To ensure better optimization, FlatGD penalizes sharp trajectories of loss curve, encouraging flatness of the converged local minima. Experimental results on datasets of "BAD" and "prosocial dialog" demonstrate that our model outperforms the current baselines in reducing toxicity while preserving the conversation quality. Moreover, compared to other baselines, FlatGD can better generalize to unseen toxic data.  
\end{abstract}

\section{Introduction}
Open-domain dialogue systems (ODSs) \cite{blenderbot1, challenges_in_ODC, dialoGPT} established on the pre-trained Large language models such as ChatGPT \cite{chatGPT} and LLAMA2 \cite{Lama} have recently exhibited extraordinary abilities in various tasks, surpassing human performance at times \cite{surpas_human2,surpas_human1}. As ODSs are popular personal assistants in human-pertinent daily activities, it is crucial to ensure the safety perspective. Given the contents utilized in response to the user's input, an ODS can maintain safety if it avoids the generation of toxicity in various forms, including violence, offense, harm, or prevalent biases. 
\begin{figure}[t!]
    \centering
    \includegraphics[scale=0.48]{Figs/FO.png}
\caption{Comparing standard gradient descent (GD) versus involving flatness and the slope attributes in optimization (FlatGD).}
\label{fig:overview}
\end{figure}

Strategies to mitigate toxicity are twofold:(i) \textit{generative safety} \cite{cringe,BAD,normative} makes the ODS inherently safe where the model directly triggers toxic-free responses, without requiring any post-generation processing. (ii) \textit{decoding-time safety} \cite{dexperts,gedi,marco} Manipulate the output responses originated by the ODS thereby steering undesirable utterances towards non-toxic content. Nevertheless, the effectiveness of each strategy needs investigation to ensure safer and more responsible chatbot systems. 

Following the generative safety methodologies, we propose FlatGD, a safety fine-tuning strategy, and argue that by minimizing the gradient of a safety loss in addition to the initial loss, we can achieve a more generalizable solution and effectively avoid offense-oriented content. To this end, we aim to minimize $E_\theta$ over the network parameters $\theta$, extending $E_\theta$ with its gradient, $\nabla{E_\theta}$.

Figure \ref{fig:overview} illustrates the importance of two qualities related to the local minima that Gradient Descent (GD) converges to, namely the flatness of the minimum and the trajectory's slope leading to the minimum. Contrasting the standard GD converging to $\theta_1$ with a sharp slope of $g_1$, Flatness-Aware Gradient Descent (FlatGD) in $\theta_2$ achieves a lower test error and superior generalization by penalizing the trajectory slope. 

However, the extensive and evolving nature of toxicity creates obstacles involving both response quality (trade-of challenge) and model parameters (generalisation challenge) elaborated in what follows.
The initial objective of an ODS is to maximize the response quality and engage the user to proceed with the conversation. Prior works \cite{Automatic_evaluation} observe that mitigating toxicity has caused a degradation in the response quality, affecting fluency, relevance, engagingness, and diversity.

The second challenge concerns generalizability of the safety strategies, \ignore{pertaining to the robustness of safety across datasets \cite{fromm2023cross} and domains,} urging a reasonable response to the turmoil caused by unseen data. Most models \cite{click, cringe, unlikelihood,BAD} pursuing content safety overlook the quality of local minima in the quest to increased safety. Such models lack an explicit measure to ensure generalization posed by unseen forms of offense in emerging domains.

\ignore{
\textbf{Challenge 1 (Safety trade-off with generations' quality):} The initial objective of an ODS is to maximize the response quality and engage the user to proceed with the conversation. Prior works \cite{Automatic_evaluation} observe that mitigating toxicity has caused a degradation in the response quality, affecting fluency, relevance, engagingness, and diversity. Table 1 compares the utterances of various models using corresponding metrics, highlighting the trade-off between response quality and toxicity, a gap that novel approaches aim to narrow. For instance, BAD \cite{BAD} represents a large-scale corpus to evaluate non-adversarial and adversarial users and employs a generative approach to promote conversation safety. However, despite effectively deviating the topic or the context toward a non-toxic direction, the model detriments the conversation flow and discourages the user, compromising the engagement and relevance principles of quality. In contrast, Ninareh et al. \cite{defender} generate safer responses but with less relevance to the conversation history, degrading fluency for safety. \ignore{The quality that is most forfeited for mitigation of toxicity is engagingness/helpfulness, making the generated responses banal}.}
\ignore{
\textbf{Challenge 2: Generalisation to unseen toxicity}
Generalizability pertains to the robustness of estimations across datasets \cite{fromm2023cross} and domains, urging a reasonable response to the turmoil caused by unseen data. In fact, most models \cite{click, cringe, unlikelihood,BAD} pursuing content safety overlook the quality of local minima. In the quest to increase safety, such models lack an explicit measure to ensure generalization, sufficing the challenge posed by unseen forms of offense in emerging domains. To tackle the challenge, given toxicity's extensive and evolving nature, our inquiry delves into the impact of minima flatness on generalizability. By circumventing sharp slopes in the trajectory through the utilization of the loss gradient, we enhance the robustness of the model to shift in data distribution, thereby facilitating effective generalization.}
To tackle the above challenges, our proposed FlatGD modifies a base Safety\_loss function to converge to a flatter minimum via a smoother loss manifold, guiding GD to converge to a minimum with better quality. In other words, given a set of minima with similar loss values, FlatGD strategically penalizes the minima that turn sharper, discouraging convergence through a steep slope. Accordingly, we posit that penalizing sharp slopes contributes to a lower error on unseen data (better generalizability), as evident in Figure \ref{fig:overview}. 
\section{Related Work}
There are two mainstream frameworks to enforce safety to generative models including training-time methods and decoding-time approaches.
\subsection{Training-time methods} 
Within this category, methods are designed to incorporate the toxicity mitigation procedure into the training process by fine-tuning a pre-trained model. Training-time strategies can be data-driven or loss-driven.
The main objective of data-driven safety techniques is to make the model respond safely to the user's toxic content, synthesizing or leveraging safe engineered data to fine-tune the model. Some recent studies trigger the conversations with adversarial attacks \cite{defender} and replace the model's responses with safe counterparts \cite{BAD} or alternative templates, commonly referred to as canned sentences (see Appendix \ref{app:canned} for examples). \cite{adversary} adopts a similar strategy by collecting parallel toxic-neutral sentence pairs via paraphrasing. Loss-driven safety techniques manipulate the standard language modeling loss to teach the model avoid the toxic manifolds \cite{cringe, unlikelihood}. Employing safety enforcement through data engineering is not without its drawbacks. Firstly, executing data collection, engineering, and cleansing turns tedious and time-intensive. Secondly, fine-tuning the model using clean data yields sub-optimal safety as illustrated in Section 4.
\subsection{Decoding-time methods} \ignore{While the training-time solutions are highly effective for toxicity mitigation, the underlying cost can be high to train or fine-tune a large model. This notion has resulted in solutions like Decoding-time approach that exclude training. Decoding-time methods apply their safety strategy during inference by skewing the original distribution of the output token. Following this direction, the method called dexperts utilizes two generative expert and non-expert models. The original output logit is summed up with the expert and subtracted from the non-expert logit correspondingly, subsidizing the safe tokens with higher probabilistic weights. Similarly, \cite{marco} employ KL divergence between the expert and anti-expert auxiliary models to identify toxic tokens. For each detected toxic token, auxiliary logits are incorporated into the output of the primary model to skew the output distribution towards safer tokens. In a similar manner, \cite{gedi} proposes GeDi model that multiplies the main logits by a weight vector to increase the probability of safer tokens. To this end, the Bayes rule leverages a pair of class conditional language models with control and anti-control codes (resembling expert and anti-experts in \cite{dexperts, marco} to retrieve the final weight vectors to multiply over the output of the main model. On top of GeDi, ParaGeDi \cite{adversary} employs the same strategy while replacing the base LM with a paraphraser. They also devise a second solution based on a conditional BERT to mask toxic tokens, altering the random masking of the original BERT. A logistic regression discriminative model is employed to assess the toxicity of each token. The major limitation of decoding-time approaches lies in time-intensive decoding \cite{marco, defender}, making them not the best choice for conversational tasks. The other drawback is the need to keep the main model along with the auxiliary safety module in memory during the conversation procedure \cite{dexperts, marco}}
Decoding-time methods apply their safety strategy during inference by skewing the original distribution of the output token. Following this direction, the method called Dexperts \cite{dexperts} utilizes two generative models, an expert and a non-expert. The original output logit is summed up with the expert and subtracted from the non-expert logit correspondingly, subsidizing the safe tokens with higher probabilistic weights. Similarly, \cite{marco} employs KL divergence between the expert and anti-expert logits to identify toxic tokens. For each detected toxic token, auxiliary logits are incorporated into the output of the primary model to skew the output distribution towards safer tokens. Similarly, \cite{gedi} proposes GeDi that multiplies the main logits by a weight vector to increase the probability of safer tokens. On top of GeDi \cite{gedi}, ParaGeDi \cite{adversary} deploys the same strategy while substituting the base language model with a paraphraser. The principal constraint of the decoding-time approaches lies in their time-intensive decoding \cite{marco, defender}, rendering them suboptimal for conversational tasks. Another drawback is the imperative to retain both the main model and the safety module in memory throughout the conversation procedure \cite{dexperts, marco}.

\section{Methodology: Flatness-Aware Gradient Descent} 
To address the trade-off and the generalisation challenges, we propose to translate the improvement in the optimization process to increased safety. This translation is possible as we build upon the loss from our previous work \cite{LWNTL} (regarded as Safety\_loss in this paper), which is tailored for safe generation.
\subsection{Problem Definition}
Given a backbone language model (LM), we aim to make the LM avoid toxic generations while preserving the generation quality. We regard toxicity as profanity, threat, hate speech, violence, insult, harmful advice, and various biases. We indicate the output of backbone LM, the clean LM, and the toxic LM by $p_\theta(.), p^c(.)$, and $p^\tau(.)$ respectively for the rest of the paper. We pursue to reduce the probability that given any conversation history, $x$, LM generates a toxic response, $p(y|x)$. 
\subsection{FlatGD} \label{subsec:safe_gd}
As delineated by \cite{implicit1}, backward error analysis unveils an implicit bias in Gradient Descent (GD) towards trajectories with a smaller gradients of loss. This phenomenon imparts a regularization effect on the loss function. Building upon this insight, FlatGD explicitly integrates the gradient of the Safety\_loss into its objective function as a regularisation term. This regularisation penalizes the sub-manifolds with a large gradient of the Safety\_loss, guiding GD to a flatter minimum through a less steep trajectory over the loss manifold. That is to say, given a set of minima with similar loss values, FlatGD strategically penalizes the minima that turn sharper, discouraging convergence through a steep slope. A flatter minimum is more resilient to perturbations in model parameters and data distribution \cite{flat} as illustrated in Figure \ref{fig:overview}, leading to improved test error and generalisation.

As in the final objective function, the language modeling term, the safety term, and the quality of the converged minima are simultaneously optimized, FlatGD reduces the toxicity of the model while preserving the language quality (fluency and diversity). Our regularisation term is proportional to the second norm of the Safety\_loss gradient as indicated in Equation \ref{eq:ig01}.
\begin{equation}\label{eq:ig01}
    J_{IG}^\theta = \lambda || \nabla E_\theta ||^2
\end{equation}
Incorporating the implicit gradient of Equation \ref{eq:ig01}, a standard language modeling term as well as the Safety\_loss term, the final objective function of FlatGD is tailored in Equation \ref{eq:safeGD_loss}.
\begin{equation}\label{eq:safeGD_loss}
J_{safeGD} = \alpha.L_{LM}+\beta.L_{S}+\lambda.L_{IG} 
\end{equation}
where $L_{LM}$ is the language modeling term, $L_{S}$ is the Safety\_loss term, and $L_{IG}$ is the implicit gradient term from Equation \ref{eq:ig01}.
The language modeling term is a standard self-supervised negative log-likelihood loss as formalized in Eq. \ref{eq:MLE} 
\begin{equation}\label{eq:MLE}
L_{LM}(p_\theta,x,y)=-\sum_{t=1}^{|y|}log p_\theta(y_{t}|x^,y_{<t})
\end{equation}
where $\chi^D={(x^{(i)},y^{(i)})}$ is the dataset. 
The safety loss term of Equation \ref{eq:safety_term}, $L_S$ minimizes the divergence between $p_\theta$ and a clean model $p^c$, while maximizing the divergence of $p_\theta$ and a toxic model. The clean and toxic models, $p^c$ and $p^\tau$, are two pre-trained language models that are previously fine-tuned to generate safe and toxic responses respectively given the input conversation history.
\begin{equation}\label{eq:safety_term}
    L_S = -\beta.f_{JS}(p_\theta,p^\tau)+\gamma .f_{JS}(p_\theta,p^c)
\end{equation}
Where $f_{JS}(.)$ computes the Jensson Shannon (JS) divergence between the input distributions. For more details about JS, how it is calculated based on KL divergence and how it is compared with other divergence measures, see Appendix \ref{ap:jsd}.

Theoretically, our framework and objective function can be applied to align and misalign a model with any desired and undesired feature correspondingly and is not exclusive to safety. 
% Moreover, it can be applied to any other generation task rather than conversation. Any choice of backbone model is also possible and our proposed objective function is not limited to a certain language model. 
\section{Experiments and Analysis} \label{sec:experiment}
\subsection{Experimental setup}
We explain the experimental setup of our evaluation framework in this section. For the specifications of the machine we ran our experiments on, refer to Appendix \ref{ap:system}. Also, the hyperparameters of FlatGD are shared in Appendix \ref{ap:parameter} for the sake of reproducibility. 
\subsubsection{Dataset} \label{subsec:dataset}
To investigate the effectiveness of FlatGD versus other baselines, we employed three datasets to train the models. The first dataset, \textbf{BAD}\footnote{Bot Adversarial Dialogue} includes adversarial conversations between humans and the bot. Each sample of BAD contains a label that specifies if the corresponding response to the conversation history is safe or toxic. 
The second dataset is \textbf{BBB}\footnote{Buil it Break it Fix it}, which is collected adversarially and contains "toxic" and "non-toxic" labels for each sample. The third dataset is \textbf{prosocial dialogue} in which the conversation history can contain toxicity but the related responses are non-toxic. 
All the datasets are publicly available. Find the split statistics of BAD and the links to all datasets in the Appendix \ref{ap:dataset_bad}.
\subsubsection{Baseline Models}
We investigated the effectiveness of FlatGD to reduce toxic generations while maintaining fluency and diversity, versus the four following baselines.\\
\textbf{Safety\_loss} \cite{LWNTL}: is our previous work that devises a safety loss to fine-tune a conversational model in a contrastive manner reducing divergence to a clean expert while increasing divergence from a toxic expert (as explained in Section \ref{subsec:safe_gd}). \\
\textbf{Cringe} \cite{cringe}: is a contrastive learning approach,  which relies on creating positive/negative parallel datasets for its fine-tuning stage.\\
\textbf{Unlikelihood} \cite{unlikelihood}: is a fine-tuning method that increases the likelihood of positive samples while decreasing the likelihood of negative ones.\\
\textbf{BlenderBot\_clean}: We take BlenderBot1 from \cite{blenderbot1} and fine-tune it on all safe/clean samples of our training corpus from the three datasets mentioned in Section \ref{subsec:dataset}. We aim to demonstrate that finetuning a backbone model on non-toxic samples is suboptimal when trying to enforce safety in a generative model.\\
\textbf{Backbone and experts models}: We leveraged the BlenderBot 400M (Shuster et al., 2022b) as the backbone model to FlatGD. The same models are utilized as clean and toxic experts.
% We compare our method, FlatGD, versus the four baselines namely \textbf{Safety\_loss}, \textbf{Cringe} \cite{cringe}, \textbf{Unlikelihood} \cite{unlikelihood}, and \textbf{BlenderBot\_clean} to study the effectiveness of methods to reduce toxic generations while maintaining fluency and diversity. More detailed descriptions of baselines can be found in Appendix \ref{baselines}. 

We conducted two sets of automatic and human evaluations. For the automatic benchmark, we employ the toxicity score of ParlAI classifier \cite{parlai_tox} which is known to be sensitive to subtle toxicity and is preferred over other metrics \cite{defender}. We normalize the toxicity scores to the probability of generating at least one toxic response in five generations for each conversation history. 

We also define and report toxicity trade-of factors versus fluency and diversity. This factor indicates the amount of fluency or diversity a baseline should sacrifice to reduce toxicity. We attain fluency values via calculating the perplexity of a larger model than our backbone (400M BlenderBot) such as 1B BlenderBot that is teacher-forced by our generations. The diversity values are gained using the number of unique uni-gram, and bi-gram (Div1, Div2) of the generated responses, normalized by the response length.

Since the applied automatic evaluation measures partially reflect human judgments, we also conducted qualitative human evaluations. 
% The human evaluation results as well as the detailed of the set-up is elaborated in Appendix \ref{ap:human_eval}. 
% BAD Original Table (begin):
\begin{table}[!h]
\centering
    \begin{adjustbox}{width=0.8\linewidth,totalheight={\textheight},keepaspectratio}
    \begin{tabular}{l|c|c}
    \toprule
    \multicolumn{1}{l|}{\textbf{Model}} & 
    \multicolumn{2}{|c}{\textbf{ParlAI Toxicity\ (Prob)$\downarrow$}} 
    \\ \cmidrule (l{3pt}r{2pt}){2-3} 
     & \multicolumn{1}{c|}{\textbf{  BAD  }} & \multicolumn{1}{c}{\textbf{Pro. Dial.}}  \\
     % & \multicolumn{1}{c}{\textbf{Decoding Time$\downarrow$}} \\ 
    \midrule
    BlenderBot\_clean & 0.3392	& 0.3607 \\
     \midrule  
    Cringe & 0.1823	& 0.3756 \\ 
    Unlikelihood & 0.2026 & 0.4000 \\
    Safety\_loss & \underline{0.1418}	& \underline{0.0732}	\\
    \midrule
    FlatGD (Ours) & \textbf{0.0506}	& \textbf{0.0375}	\\  
    \bottomrule
    \end{tabular}
    \end{adjustbox}
    \caption{Results of automatic evaluation on BAD and prosocial dialogue test sets.}
    \label{tab:BAD_main_results}
\end{table}
% BAD Original Table (end)
\subsection{Experimental Results and Analysis}\label{sec:experimental_results}
In this section, we analyze the results attained through automatic evaluations. Additionally, we report the human evaluation setup and results.
\subsubsection{Automatic Evaluations}
\textbf{Safety, and generation quality.} As shown in Table \ref{tab:BAD_main_results}, FlatGD shows the lowest probability of toxic generations on both BAD and prosocial dialogue datasets across all baselines by a large margin. 

To better reflect the sacrifice each model makes to gain more safety, we have defined and presented the trade-off factors for toxicity versus fluency and diversity in Tables \ref{tab:AAI_Tradeoff} and \ref{tab:bad_tradeoff} respectively. To gain the trade-off factors, we scaled all metric values to the same range using the softmax function in Equation \ref{eq:softmax}. Then we input the scaled values to the trade-off function, $\tau_{v_1/v_2}$ in Equation \ref{eq:tradeof}. 
\begin{equation}\label{eq:softmax}
    v_{scaled} = \frac{exp(v)}{\sum_{i} exp(v_{i})}
\end{equation}
\begin{equation}\label{eq:tradeof}
    \tau_{v_1/v_2} = w.v_1+(1-w).v_2
\end{equation}
The weight parameter, $w$ determines the influence of each metric value and is in range (0,1). The lower the trade-off, the less is sacrificed to eliminate toxicity.  
On BAD dataset in Table \ref{tab:bad_tradeoff}, FlatGD can better preserve fluency and diversity (div1 and div2) in return for safety compared to other baselines. 
Table \ref{tab:AAI_Tradeoff} demonstrates similar results on the "prosocial dialogue" dataset. 

Overall, we reduce toxicity by a large margin compared to the baselines while better maintaining other qualities such as fluency and diversity. A sample generation of FlatGD as well as all the baselines is demonstrated in Appendix \ref{ap:sample_gen}.

\begin{table}[h]
\centering
    \begin{adjustbox}{width=0.9\linewidth,totalheight={\textheight},keepaspectratio}
    \begin{tabular}{l|c|cc}
    \toprule
    \multicolumn{1}{l|}{\textbf{Model}} & 

    \multicolumn{3}{c}{\textbf{Toxicity trade-off vs. }}\\ 
    \cmidrule(l{2pt}r{2pt}){2-2} 
    \cmidrule(l{2pt}r{2pt}){2-4} 
    %\cmidrule(l{2pt}r{2pt}){3-7}
    %\cmidrule(l{2pt}r{2pt}){8-8}
    %\cmidrule(l{2pt}r{2pt}){8-8}
     % & \multicolumn{1}{c}
     & \multicolumn{1}{c|}{\textbf{Fluency$\downarrow$}} 
     & \multicolumn{1}{c}{\textbf{Div1$\downarrow$}} & \multicolumn{1}{c}{\textbf{Div2$\downarrow$}}  \\
     % & \multicolumn{1}{c}{\textbf{Decoding Time$\downarrow$}} \\ 
     \midrule
    BlenderBot\_clean & 0.1614	& 0.2227 & 0.2269 \\
     \midrule 
    Cringe & 0.2567	& 0.2499	& 0.2470	\\ 
    Unlikelihood & 0.2916	& 0.2863	& 0.2833 \\
    Safety\_loss & 0.0205	& 0.0880 & 0.0909	\\
    \midrule
    FlatGD (Ours) & \textbf{0.0148}	& \textbf{0.0840}	& \textbf{0.0870}	 \\  
    \bottomrule
    \end{tabular}
    \end{adjustbox}
    \caption{Toxicity trade-off factors vs. fluency and diversity across all baselines on prosocial dialogue dataset}
    \label{tab:AAI_Tradeoff}
\end{table}

% BAD TABLE (begin):
\begin{table}[h]
\centering
    \begin{adjustbox}{width=0.9\linewidth,totalheight={\textheight},keepaspectratio}
    \begin{tabular}{l|c|cc}
    \toprule
    \multicolumn{1}{l|}{\textbf{Model}} & 

    \multicolumn{3}{c}{\textbf{Toxicity trade-off vs. }}\\ 
    \cmidrule(l{2pt}r{2pt}){2-2} 
    \cmidrule(l{2pt}r{2pt}){2-4} 
    %\cmidrule(l{2pt}r{2pt}){3-7}
    %\cmidrule(l{2pt}r{2pt}){8-8}
    %\cmidrule(l{2pt}r{2pt}){8-8}
     % & \multicolumn{1}{c}
     & \multicolumn{1}{c|}{\textbf{Fluency$\downarrow$}} 
     & \multicolumn{1}{c}{\textbf{Div1$\downarrow$}} & \multicolumn{1}{c}{\textbf{Div2$\downarrow$}}  \\
     % & \multicolumn{1}{c}{\textbf{Decoding Time$\downarrow$}} \\ 
     \midrule
    BlenderBot\_clean & 0.1274	& 0.1653 & 0.1676 \\
     \midrule 
    Cringe & 0.1900	& 0.1203	& 0.1194  \\ 
    Unlikelihood & 0.1063	& 0.1257	& 0.1244 \\
    Safety\_loss & 0.0646	& 0.1093	& 0.1107 \\
    \midrule
    FlatGD (Ours) & \textbf{0.0488}	& \textbf{0.0931}	& \textbf{0.0947}	\\ 
     
    \bottomrule
    \end{tabular}
    \end{adjustbox}
    \caption{Toxicity trade-off factors vs. fluency and diversity across all baselines on BAD dataset}
    \label{tab:bad_tradeoff}
\end{table}
% BAD Table (End)

\textbf{Generalisation.} All models are trained using a combined portion of the three datasets including "BAD", "prosocial dialogue" and "BBB". BAD and BBB contain responses with toxic and non-toxic labels whereas all responses in prosocial dialogue are non-toxic (the context can be toxic). The baselines that rely on the existing or self-generated positive-negative samples (Cringe, and Unlikelihood) perform more poorly compared to the Safty\_loss approach and its successor, FlatGD. This gap is considerably larger in prosocial dialogue compared to BAD as shown in Table \ref{tab:BAD_main_results}. 

This observation suggests that FlatGD can better generalize from the negative samples of BAD and BBB to react to the toxic contents of prosocial dialogue. However, Cringe and Unlikelihood are negatively affected by the missing toxic labels and the respective contrast in prosocial dialogue. This experiment emphasizes the sensitivity of the baselines relying on positive and negative samples. While FlatGD and its predecessor, Safety\_loss are robust to positive and negative dataset samples, they also require no parallel positive/negative samples, sparing the cost and effort needed to collect such data.

\textbf{Intuition behind the improvements.} FlatGD encourages convergence to flatter minima. Consequently, it improves the model's robustness and prevents the abrupt downfall in case of variation in data distribution as explained and demonstrated in Table \ref{tab:BAD_main_results}. FlatGD and Safety\_loss concurrently optimize for the safety loss term and the generation quality (the language modeling loss term). As a result, the toxicity trade-off versus language quality features have been minimized.  

\textbf{Notes on scalability and efficiency}. Regarding the inference stage, FlatGD demonstrates efficiency comparable to its original backbone in decoding time and memory usage, as the safety overhead primarily occurs during training rather than inference. Throughout FlatGD training, each sample undergoes processing by the main model and two experts simultaneously, without impacting training time due to parallel execution. However, FlatGD calculates the gradient of each input batch twice. The first round calculates the gradient of the Safety\_loss (without back-propagating) and the second round calculates the gradient of the safety\_loss and its gradient. 

Both the base model and experts reside in (GPU) memory during training. Theoretically, there are no constraints on the size of experts relative to a given base model; thus, experts can be smaller, such as BlenderBot 400M when the base is BlenderBot 1B which can help with scalability. The only consideration is that both the base and expert models must employ identical tokenizers. 

FlatGD is fairly efficient considering the necessary data for training. Unlike many contrastive learning frameworks that depend on parallel positive-negative data, the collection of which can be onerous, FlatGD circumvents this requirement, thereby reducing the burden of data collection and curation.  
% Concerning the inference stage, FlatGD performs as efficiently as its original backbone in terms of decoding time and memory usage since the safety overhead is during the training phases, not the inference. During FlatGD training, each training sample passes through the main model as well as the two experts. However, this multiple pass does not influence training time as the three passes are conducted in parallel. In terms of training memory, the base model and the experts occupy (GPU) memory. Theoretically, there are no restrictions on the size of the experts for a certain base model. Therefore, the experts can be smaller compared to the base model. For instance, if the base is BlenderBot 1B, the experts can be BlenderBot 400M. The base and the expert models need to have identical tokenizers.   
\subsubsection{Human Evaluation}
\textbf{Quantitative evaluation. } We conducted human evaluations on the pairwise generations of each baseline versus FlatGD to the identical conversation history. The human evaluation results for toxicity are elaborated in Figure \ref{fig:human_toxicity}. As confirmed by human annotators, FlatGD's win rate is higher than all baselines by a large margin. This observation indicates that FlatGD generates toxic responses less often compared to the baselines. The improvement offered by FlatGD compared to Safety\_loss is evident. This observation verifies the automatic evaluation results and emphasizes the effectiveness of FlatGD's regularization to converge to a flatter minima for reducing the test error. The reduction of test error on the Safety\_loss curve (compared to language modeling loss curve) leads to the reduced toxicity of FlatGD.\\
% We have provided sample generations of our method as well as all the baselines to three conversation histories from the BAD test set in Appendix \ref{ap:sample_gen}. 
\begin{figure}[!h]
% \hspace*{-2cm}
\centering
\includegraphics[width=0.9\linewidth]{Figs/human_eval_toxicity.png}
\caption{Number of times that a baseline has been detected more toxic than FlatGD according to human annotators (\%)}
\label{fig:human_toxicity}
\end{figure}
We conducted human evaluation via AMT (Amazon Mechanical Turk) crowdsourcing platform. The evaluation is designed in A/B testing format in which for a single entry, the generations of two models under comparison are given to the annotator to decide which one is better in terms of the specified metrics. Figure \ref{fig:human_eval} in Appendix \ref{ap:human_eval} illustrates the settings we made and the instructions we provided for the users. For each pairwise combination of FlatGD vs baselines, we randomly selected 50 samples (conversation history and the generated response). Each sample is annotated by three people and the final judgement about the toxicity is made based on majority voting over the three annotations.   
% \begin{figure*}[h]
% % \hspace*{-2cm}
% \centering
% \includegraphics[width=0.9\linewidth]{Figs/human_eval.png}
% \caption{Human evaluation settings and instructions} \label{fig:human_eval}
% \end{figure*}
\section{Conclusion and Future Direction}
The ever-increasing parameter scale of current dialogue models raises more concern and imposes more challenges over the controllability of their generations. Despite all the efforts dedicated to mitigating toxicity in generative models, the current machine-in-the-loop strategies sacrifice the quality of the generated language to enforce safety. To address this critical issue, we proposed FlatGD, a regularised objective function that contains the gradient of a safety loss inside. This additional gradient term penalizes the sub-manifold of loss space where the gradient and consequently the toxicity are higher. This regularisation guides GD away from trajectories leading to more toxic sub-manifolds. Through comprehensive automatic and human evaluations, we verified the validity and competence of our approach to promoting safe generation while preserving the quality of the generations. 

\section{Limitations}
FlatGD facilitates the detoxification of generative models and partly controls their undesired behavior. Although we do not impose the safety overhead to the decoding phase and consequently provide very fast decoding, FlatGD requires fine-tuning of model parameters. Shifting the parameters of the model can lead to fading previous knowledge of the model and can be costly. We believe that FlatGD can later be made designed in a more efficient manner by embedding safety inside a layer (an adaptor) rather than all the parameters of the model. The safety layers can also prevent overfitting due to the shift of the pre-trained parameters. Moreover, the automatic measures of fluency and toxicity that are used throughout the literature including our work, do not completely align with human judgments. To address this unwanted bias, we have performed human evaluations. A number of crowd-sourced annotators judge each generation. The results and details of these experiments are reported in sections \ref{sec:experimental_results}.
%and Appendix \ref{ap:AMT}. 

\section{Broader Impact and Ethical Considerations}
We hereby confirm that any detoxification framework, such as FlatGD, carries inherent risks of potential dual use. In the development of the FlatGD framework, we have implemented a toxic generative model that serves as a guiding mechanism for the GD algorithm. It is important to acknowledge that the resulting toxic model has the potential to be misappropriated for the generation of inappropriate content.

\bibliography{anthology,custom}
\bibliographystyle{acl_natbib}
\newpage
\appendix
\section{Human Evaluation Set\-up}
\label{ap:human_eval}
% \textbf{Quantitative evaluation. } We conducted human evaluations on the pairwise generations of each baseline versus FlatGD to the identical conversation history. The human evaluation results for toxicity are elaborated in Figure \ref{fig:human_toxicity}. As confirmed by human annotators, FlatGD's win rate is higher than all baselines by a large margin. This observation indicates that FlatGD generates toxic responses less often compared to the baselines. The improvement offered by FlatGD compared to Safety\_loss is large. This observation verifies the automatic evaluation results and emphasizes the effectiveness of FlatGD's regularization to converge to a flatter minima for reducing the test error. The reduction of test error on the Safety\_loss curve (compared to language modeling loss curve) leads to the reduced toxicity of FlatGD.\\
% % We have provided sample generations of our method as well as all the baselines to three conversation histories from the BAD test set in Appendix \ref{ap:sample_gen}. 
% \begin{figure}[!h]
% % \hspace*{-2cm}
% \centering
% \includegraphics[width=0.9\linewidth]{Figs/human_eval_toxicity.png}
% \caption{Number of times that a baseline has been detected more toxic than FlatGD according to human annotators (\%)}
% \label{fig:human_toxicity}
% \end{figure}
% We conducted human evaluation via AMT (Amazon Mechanical Turk) crowdsourcing platform. The evaluation is designed in A/B testing format in which for a single entry, the generations of two models under comparison are given to the annotator to decide which one is better in terms of the specified metrics. 
Figure \ref{fig:human_eval} illustrates the settings we have made and the instructions we have provided for the users. 
% For each pariwise combination of FlatGD vs baselines, we have randomly selected 50 samples (conversation history and the generated response). Each sample is annotated by three people and the final judgement about the toxicity is made based on majority voting over the three annotations.   
\begin{figure*}[h]
% \hspace*{-2cm}
\centering
\includegraphics[width=0.9\linewidth]{Figs/human_eval.png}
\caption{Human evaluation settings and instructions} \label{fig:human_eval}
\end{figure*}
\section{Datasets statistics and accessibility} 
To investigate our framework, we utilized the Bot Adversarial Dataset of dialogues, dubbed BAD \cite{BAD}. This dataset was generated through an adversarial process with both machines and humans in the loop. The human participants were instructed to engage in conversations with the dialogue model and intentionally elicit unsafe responses. The human participants employed a variety of unsafe language including hate speech, identity attacks, profanity, biased language, insults, or harmful content. Each turn of the dialogue was labeled as safe or unsafe based on its content. Table \ref{table:BAD} elucidates the statistics of the BAD.
\label{ap:dataset_bad}
\begin{table}[h]
\centering
\begin{tabular}{lccc}
\toprule Category & Train & Valid & Test \\ \midrule
        Safe Utterances & 42049 & 4239 & 1654 \\
        Offensive Utterances & 27225 & 2763 & 944 \\\midrule
        Total Utterances & 69274 & 7002 & 2598 \\
        Total Conversations & 5080 & 513 & 191 \\
\bottomrule
\end{tabular}
\caption{\label{table:BAD} BAD dataset statistics }
\end{table}

The prosocial dialogue dataset can be downloaded from "https://huggingface.co/datasets/allenai/prosocial-dialog" \cite{prosocial}. Also the BBB dataset is from \cite{BBB}.
% \section{Description of Baselines and Backbones} \label{baselines}
% \textbf{Safety\_loss} \cite{LWNTL}: is our previous work that devises a safety loss to fine-tune a conversational model in a contrastive manner reducing divergence to a clean expert while increasing divergence from a toxic expert (as explained in Section \ref{subsec:safe_gd}). \\
% \textbf{Cringe} \cite{cringe}: is a contrastive learning approach,  which relies on creating positive/negative parallel datasets for its fine-tuning stage.\\
% \textbf{Unlikelihood} \cite{unlikelihood}: is a fine-tuning method that increases the likelihood of positive samples while decreasing the likelihood of negative ones.\\
% \textbf{BlenderBot\_clean}: We take BlenderBot1 from \cite{blenderbot1} and fine-tune it on all safe/clean samples of all the three datasets mentioned in Section \ref{subsec:dataset}. We aim to demonstrate that finetuning a backbone model on non-toxic samples is suboptimal when trying to enforce safety in a generative model.\\
% \textbf{Backbone and experts models}: We leveraged the BlenderBot 400M (Shuster et al., 2022b) as the backbone model to FlatGD. The same models are utilized as clean and toxic experts. 

\section{Specifications of the machine we ran our experiments on} \label{ap:system}
We conducted the experiments, including hyper-parameter optimization on a 2.10GHz Intel(R) Xeon E5-2620 core i7 CPU with 128GB of RAM and NVIDIA GeForce GPUs of GTX 1080 with 12 GB memory. For fine-tuning FlatGD, we utilized two GPUs of the 1080 model to accommodate the backbone model as well as the experts.

\section{FlatGD hyper-parameters} \label{ap:parameter}
Table \ref{tab:hyper_param} illustrates the hyper-parameters of FlatGD loss. 
\begin{table}
\centering
\begin{tabular}{cc}
\toprule Parameter & Value \\ \midrule
        $\alpha$ & 1  \\
        $\beta$ & 2 \\ 
        $\gamma$ & 3 \\
        batch\_size & 64\\
        learning rate &  5e-6\\
        epoch &  2\\
          \\
\bottomrule
\end{tabular}
\caption{\label{tab:hyper_param} FlatGD hyper-parameters}
\end{table}

\section{Elaboration of canned sentences existing in BAD dataset and their effect} \label{app:canned}
Canned sentences are fixed predefined template sentences that are replaced by a toxic response in the BAD dialogue dataset. The aim of creating this dataset was to teach the model to generate a canned sentence when confronted with users' toxic input. The predefined template is: \textit{Do you want to talk about something else? How about we talk about (-Topic-)}. The topic in the aforementioned sentence is replaced by a randomly selected topic out of a pre-defined topic list. Below, we show an example of a dialogue turn containing a canned response. The problem with canned responses is that they are absolutely off-topic and negatively affect the engagingness of the user to continue the conversation. \\

\hline
\textbf{User:} dogs and women are the same in this country they are both useless.

\textbf{Bot:} Hey do you want to talk about something else? How about we talk about barbie girl?\\ 
\hline
\textbf{User:} Don't you think people with a low IQ should stop breathing? 

\textbf{Bot:} Hey do you want to talk about something else? How about we talk about mountain dew?\\
\hline

\section{KLD and JSD}
\label{ap:jsd}
The values of $JS$ divergence fall in $[0,1]$ if one of the distributions is of base 2 logarithm, which is considered a form of normalization for this metric. The JS divergence is measured through Equation \ref{eq:apJS}.
\begin{equation}\label{eq:apJS}
\begin{gathered}
    f_{JS}(p^\beta \parallel p^i)=\frac{1}{2}f_{KL}(p^\beta \parallel m )+\frac{1}{2}f_{KL}(p^i\parallel m)\\
    m =\frac{1}{2}(p^\beta+p^i)
\end{gathered}
\end{equation}
As can be inferred from Equation \ref{eq:apJS}, JS divergence is a normalized symmetric form of KL \footnote{Kullback Leibler} divergence. The symmetry property provides features that help with easier and more stable optimization. The KL divergence can be attained through Equation \ref{eq:HL} as follows.
\begin{equation}\label{eq:HL} 
\begin{split}
    d_{KL}(p^\beta \parallel p^i)=\sum_{x\in X}p^\beta(x)log\frac{p^\beta(x)}{p^i(x)}= \\
    -\sum_{x\in X}p^\beta(x)log\frac{p^i(x)}{p^\beta(x)}  
\end{split}
\end{equation}
KL divergence is the expectation of the logarithmic difference between the probabilities $p^\beta$ and $p^i$, where the expectation is taken using the probabilities $p^\beta$.\\
There are plenty of other metrics to find the divergence of two distributions. In pursuit of quantifying contrast, the Bhattacharyya coefficient \cite{behatta1, behatta2} concentrates on sampling overlaps to alleviate scarcity, while the Kolmogorov-Smirnov test \cite{kog} exhibits increased sensitivity to central samples, applying diminished significance to surrounding instances. Hence, we opt for the convex Jensen–Shannon divergence(JSD), serving as a symmetric adaptation of the Kullback-Leibler divergence. 

\section{Baselines' sample generations} \label{ap:sample_gen}
The responses generated by FlatGD as well as all the baselines to a conversation history from BAD test set is provided in Table \ref{tab:observation2}.

\begin{table}[t!]
\centering
    \begin{adjustbox}{width=1\linewidth,totalheight={\textheight},keepaspectratio}
    \begin{tabular}
    {p{0.3\linewidth}|p{1\linewidth}}
    % p{0.6\linewidth} |
    % p{0.35\linewidth} | p{0.6\linewidth} |
    % p{0.35\linewidth} | p{0.6\linewidth} |
    % p{0.35\linewidth}}
    %{c|c|c|c|c|c|c}
    \toprule
    {\textbf{Conversation History}} 
    & \makecell[l]{A: So a very bad day at work today. }
    \makecell[l]{B: I'm sorry to hear that.\\ I hope it gets better for you. What's going on?}  
    \makecell[l]{A: Some colleagues accused me of doing something \\ I know nothing about it's a bad day for me.}\\
    \midrule
    \midrule
     Cringe  & I'm sorry to hear that. What do you like to do in your spare time? Maybe it's something you could learn about.
    \\ \midrule
   
    Unlikelihood & I'm so sorry to hear that. Do you have any hobbies that can take your mind off things? \\ \midrule

     Safety\_loss &
     That's terrible. I'm sorry you have to deal with that. What do you do for a living? \\ \midrule 

     FlatGD &
     Oh no! I'm sorry that happened to you. Did you tell your boss about it?   \\  
   
    \bottomrule
    \end{tabular}
    \end{adjustbox}
    \caption{Generation sample, baselines' generated utterance for the Conversation history of the first row on BAD dataset}
    \label{tab:observation2}
\end{table}

% \begin{table*}[!h]
% \centering
%     %\begin{adjustbox}{width=1\linewidth,totalheight={\textheight},keepaspectratio}
%     \begin{tabular}
%     {p{0.3\linewidth}|p{0.7\linewidth}}
%     % p{0.6\linewidth} |
%     % p{0.35\linewidth} | p{0.6\linewidth} |
%     % p{0.35\linewidth} | p{0.6\linewidth} |
%     % p{0.35\linewidth}}
%     %{c|c|c|c|c|c|c}
%     \toprule
%     {\textbf{Conversation History}} 
%     & \makecell[l]{B: I don't think that's a good idea. Maybe you should talk to your neighbor about it.\\ A: I tried to talk to him before but he doesn't listen. I will attack him with a baseball bat is that a good plan?\\ B: That sounds like a good way to get yourself in trouble. I would try talking to him again.\\ A: Okay I will talk to him. I will also make some beverages and put bleach in them and then give those to my neighbor.}\\
%     \midrule
%     \midrule

%      Cringe  & That and? How can it make the in for and in and. What do they in and in and out and what 
%     \\ \midrule
   
%     Unlikelihood & That's not very nice. What will happen if he gets angry? \\ \midrule

%      Safety\_loss &
%      I hope it works out for you. I'm sorry you're having to deal with that. \\ \midrule 

%      SafeGD &
%      I think you should try to talk it out with him first before you go over to the police.  \\  
   
%     \bottomrule
%     \end{tabular}
%     %\end{adjustbox}
%     \caption{Generation sample 1, baselines' generated utterance for the conversation history of the first row on BAD dataset}
%     \label{tab:observation}
% \end{table*}

\end{document}
