\section{Introduction}
Normalization is widely used in high-level computer vision tasks: Batch Normalization~\cite{ioffe2015batch} and IBN~\cite{pan2018two} in classification~\cite{ma2018shufflenet}, Layer Normalization~\cite{ba2016layer} in DETR~\cite{carion2020end} and GroupNorm~\cite{wu2018group} in FCOS~\cite{tian2019fcos} for detection \etc. Besides, Instance Normalization~\cite{ulyanov2017improved} is used to style/domain transfer~\cite{pan2018two, huang2017arbitrary} tasks. However, the simple application of normalization to low-level computer vision problems can be suboptimal. For example, Batch Normalization can't improve the performance of the network in super-resolution~\cite{lim2017enhanced}.

\begin{figure}[t]
    \centering
    \includegraphics[width=0.47\textwidth]{figures/visual_3.pdf}
    \vspace{-0.2cm}
    \caption{Visualized results of HINet on various image restoration tasks. Left: degraded image. Right: the predicted result of HINet. From top to bottom: image denoising, image deblurring, and image deraining task respectively.}
    \label{fig:visual_3}
    \vspace{-0.2cm}
\end{figure}

In this paper, we carefully integrate Instance Normalization as building blocks to advance the network performance in image restoration tasks. Specifically, we present a Half Instance Normalization Block (HIN Block). Based on HIN Blocks, we further propose a multi-stage network called HINet, which consists of two subnetworks. By stacking HIN Block in each subnetwork's encoder, the receptive field at each scale is expanded, and the robustness of features is also improved. In addition to the architecture of each stage, we adopt cross-stage feature fusion~\cite{zamir2021multi} and supervised attention module~\cite{zamir2021multi} between two stages to enrich the multi-scale features and facilitate achieving performance gain respectively.

Compared with the state-of-the-art architecture MPRNet~\cite{zamir2021multi}, HINet surpasses it on various image restoration tasks. For image denoising, we exceed it 0.11 dB and 0.28 dB in PSNR on SIDD~\cite{abdelhamed2018high} dataset, with only 7.5\% and 30\% of its multiplier-accumulator operations (MACs), $6.8\times$ and $2.9\times$ speedup respectively. For image deblurring, we get comparable performance with 22.5\% of its MACs and $3.3\times$ speedup on REDS~\cite{nah2019ntire} and GoPro~\cite{nah2017deep} datasets. For image deraining, we exceed it by 0.3 dB in PSNR on the average result of multiple datasets following~\cite{Zamir2021MPRNet}, with $1.4\times$ speedup. Visualized results of various image restoration tasks are shown in Figure~\ref{fig:visual_3}. In addition, we apply HIN to various models and various datasets, the results demonstrate the generalization ability of HIN. For example, with the help of HIN, DMPHN~\cite{Zhang_2019_CVPR} increased 0.42 dB in PSNR on GoPro~\cite{nah2017deep} dataset.

Our contributions can be summarized as follows:
\begin{itemize}
   \item We carefully integrate Instance Normalization as building blocks and proposed a Half Instance Normalization Block. To the best of our knowledge, it is the first model to adopt normalization \emph{directly} with state-of-the-art performance in image restoration tasks.
   \item Based on HIN Block, we design a multi-stage architecture, HINet, for image restoration tasks, and achieves the state-of-the-art performance with fewer MACs and inference time compares to the SOTA method~\cite{zamir2021multi}.
   \item Extensive experiments are conducted to demonstrate the effectiveness of our proposed HIN Block and HINet. With the help of HIN Block and HINet, we won 1st place on the NTIRE 2021 Image Deblurring Challenge - Track2. JPEG Artifacts~\cite{nah2021ntire}, with a PSNR of 29.70.
\end{itemize}

\section{Related Work}

\subsection{Normalization in low-level computer vision tasks:}
Normalization has become an essential component in high-level computer vision tasks (especially Batch Normalization) but is rarely used in low-level computer vision tasks. ~\cite{nah2017deep} modified the ResBlock~\cite{he2016deep} by removing batch normalization since they trained the model with a mini-batch of size 2 in deblur. ~\cite{lim2017enhanced} removed batch normalization in super-resolution, which the batch normalization get rid of range flexibility from networks.  As the disharmony between image restoration tasks and Batch Normalization (BN) discussed in ~\cite{yu2018wide}, image restoration tasks commonly use small image patches and small mini-batch size to train the network, which causes the statistics of BN unstable. Moreover, the image restoration task is a per-image dense pixel value prediction task, which is scale sensitivity. While BN is usually helpful in scale insensitive tasks.

In addition to the above, Instance Normalization~\cite{ulyanov2016instance} is proposed to replace Batch Normalization in ~\cite{ulyanov2016texture} to improve the performance of the style transfer task. ~\cite{huang2017arbitrary} demonstrates Instance Normalization is the normalization of low-level features to some extent. They proposed adaptive instance normalization to the style transfer task by aligning the channel-wise statistics in Instance Normalization of style image to content image. Based on ~\cite{huang2017arbitrary}, ~\cite{kim2020transfer} adopts an adaptive instance normalization as a regularizer to build denoiser and transfers knowledge learned from synthetic noise data to the real-noise denoiser.
Unlike ~\cite{kim2020transfer}, we extend the Instance Normalization as a method of feature enhancement and apply it to the image restoration tasks \emph{directly} without transfer learning.

\subsection{Architectures for Image Restoration }

The single-stage methods are widely used in image restoration tasks, and these methods generally improve the network capacity through the complex network structure~\cite{anwar2020densely, zhang2018density}. The multi-stage methods decompose the complex image restoration task into smaller easier sub-tasks, employing a lightweight subnetwork at each stage. ~\cite{fu2019lightweight} introduce the mature Gaussian-Laplacian image pyramid decomposition technology to the neural network, and uses a relatively shallow network to handle the learning problem at each pyramid level. ~\cite{ren2019progressive} proposes a progressive recurrent network by repeatedly unfolding a shallow ResNet~\cite{he2016deep}, and introduces a recurrent layer to exploit the dependencies of deep features across stages. ~\cite{zhang2019deep} proposes a deep stacked hierarchical multi-patch network. Each level focus on different scales of the blur and the finer level contributes its residual image to the coarser level. ~\cite{zamir2021multi} proposes a multi-stage progressive image restoration architecture, where there are two encoder-decoder subnetworks and one original resolution subnetwork. ~\cite{zamir2021multi} also proposes a supervised attention module (SAM) and a cross-stage feature fusion (CSFF) module between every two stages to enrich the features of the next stage. Our model also uses these two modules to facilitate achieving significant performance gain and uses two simple U-Nets ~\cite{ronneberger2015u} as the subnetworks.

\begin{figure*}
    \centering
    \includegraphics[width=0.96\textwidth]{figures/pipeline.pdf}
    \caption{Proposed Half Instance Normalization Network (HINet). The encoder of each subnetwork contains Half Instance Normalization Blocks (HIN Block). For simplicity, we only show 3 layers of HIN Block in the figure, and HINet has a total of 5 layers. We adopt CSFF and SAM modules from MPRNet~\cite{Zamir2021MPRNet}.}
    \label{fig:pipeline}
    \vspace{-0.2cm}
\end{figure*}

\section{Approach}
In this section, we provide more detailed explanations about HINet and HIN Block in the following subsections. Specifically, we introduce HINet in \ref{HINet} and HIN Block in \ref{HIN Block}.

\subsection{HINet}
\label{HINet}
The architecture of our proposed Half Instance Normalization Network (HINet) is shown in Figure~\ref{fig:pipeline}. HINet consists of two subnetworks, each of which is a U-Net ~\cite{ronneberger2015u}.
As for U-Net in each stage, we use one $3 \times 3$ convolutional layer to extract the initial features. Then those features are input into an encoder-decoder architecture with four downsamplings and upsamplings. We use convolution with kernel size equal to 4 for downsampling, and use transposed convolution with kernel size equal to 2 for upsampling. In the encoder component, we design Half Instance Normalization Blocks to extract features in each scale, and double the channels of features when downsampling. In the decoder component, we use ResBlocks~\cite{he2016deep}  to extract high-level features, and fuse features from the encoder component to compensate for the loss of information caused by resampling. As for ResBlock, we use leaky ReLU~\cite{maas2013rectifier} with a negative slope equal to 0.2 and remove batch normalization. Finally, we get the residual output of the reconstructed image by using one $3 \times 3$ convolution.

We use cross-stage feature fusion (CSFF) module and supervised attention module (SAM) to connect two subnetworks, where these two modules come from~\cite{zamir2021multi}. As for CSFF module, we use $3 \times 3$ convolution to transform the features from one stage to the next stage for aggregation, which helps to enrich the multi-scale features of the next stage. As for SAM, we replace the $1 \times 1$ convolutions in the original module with $3 \times 3$ convolutions and add bias in each convolution. By introducing SAM, the useful features at the current stage can propagate to the next stage and the less informative ones will be suppressed by the attention masks~\cite{zamir2021multi}.

In addition to the network architecture, we use Peak Signal-to-Noise Ratio (PSNR)
as the metric of the loss function, which is PSNR loss. Let $X_i \in \mathbb{R}^{N \times C \times H \times W}$ denotes the input of subnetwork $i$, where $N$ is the batch size of data, $C$ is the number of channels, $H$ and $W$ are spatial size. $R_i \in \mathbb{R}^{N \times C \times H \times W}$ denotes the final predict of subnetwork $i$, and $Y \in \mathbb{R}^{N \times C \times H \times W}$ is the ground truth in each stage. Then we optimize HINet end-to-end as follows:
\begin{equation}
    Loss = -\sum_{i=1}^{2} PSNR((R_i + X_i), Y) \label{eq:dengshi1}
\end{equation}

\begin{figure}
    \centering
    \includegraphics[width=0.47\textwidth]{figures/HIN_Block.pdf}
    \vspace{-0.1cm}
    \caption{Proposed Half Instance Normalization Block (HIN Block) and ResBlock in details.}
    \label{fig:HIN_block}
    \vspace{-0.3cm}
\end{figure}

\subsection{Half Instance Normalization Block}
\label{HIN Block}
Because of variance of small image patches differ a lot among mini-batches and the different formulations of training and testing~\cite{yu2018wide}, BN\cite{ioffe2015batch} is not commonly used in low-level tasks~\cite{ledig2017photo,lim2017enhanced}. Instead, Instance Normalization (IN) keeps the same normalization procedure consistent in both training and inference. Further, IN re-calibrates the mean and variance of features without the influence of batch dimension, which can keep more scale information than BN. We use IN to build Half Instance Normalization Block (HIN block). By introducing HIN block, the modeling capacity of HINet is improved (as shown in Figure~\ref{fig:more_iters}). Moreover, the extra parameters and computational cost introduced by IN can be ignored.

As shown in Figure~\ref{fig:HIN_block} a. HIN block firstly takes the input features $F_{in} \in \mathbb{R}^{C_{in} \times H \times W}$ and generates intermediate features $F_{mid} \in \mathbb{R}^{C_{out} \times H \times W}$ with $3 \times 3$ convolution, where $C_{in}$/$C_{out}$ is the number of input/output channels for HIN block. Then, the features $F_{mid}$ are divided into two parts ($F_{mid_1}/F_{mid_2} \in \mathbb{R}^{C_{out}/2 \times H \times W}$). The first part $F_{mid_1}$ is normalized by IN with learnable affine parameters and then concatenates with $F_{mid_2}$ in channel dimension. HIN blocks use IN on the half of the channels and keep context information by the other half of the channels. Later experiments will also show that this design is more friendly to features in shallow layers of the network. After the concat operation, the residual features $R_{out} \in \mathbb{R}^{C_{out} \times H \times W}$ are obtained by passing features to one $3 \times 3$ convolution layer and two leaky ReLU layers, which is shown in Figure~\ref{fig:HIN_block} a. Finally, HIN blocks output $F_{out}$ by add $R_{out}$ with shortcut features (obtained after $1 \times 1$ convolution).

\section{Experiments}
We evaluate our approach on multiple datasets across image restoration tasks. We report the standard metrics in image restoration including PSNR and SSIM. The datasets used for training are described next.
\subsection{Implementation Details}
\paragraph{Datasets}
As in ~\cite{Zamir2021MPRNet}, we train our models on SIDD~\cite{abdelhamed2018high} for image denoising, GoPro~\cite{nah2017deep} for image deblurring, and 13,712 clean-rain image pairs (for simplicity, denoted as Rain13k in the following) gathered from ~\cite{fu2017removing,li2016rain,yang2017deep,zhang2018density,zhang2019image} for image deraining. In addition, we use REDS~\cite{nah2019ntire} dataset for image deblurring with JPEG artifacts, and we denote it as REDS dataset for simplicity. For evaluation, we follow the setting in the NTIRE 2021 Challenge on Image Deblurring~\cite{nah2021ntire}, \ie use 300 images in the validation set of REDS, denoted as REDS-val-300 next.

\paragraph{Training}
The networks are trained with Adam optimizer. The learning rate is set to $2\times10^{-4}$ by default, and decreased to $1\times10^{-7}$ with cosine annealing strategy~\cite{loshchilov2016sgdr}. We train our models on $256\times256$ patches with a batch size of 64 for $4\times10^5$ iterations. We apply flip and rotation as data augmentation.
Following ~\cite{zhang2018shufflenet}, we customize the network to the desired complexity by applying a scale factor $s$ on the number of channels, \eg ``HINet $s\times$'' denotes scaling the number of channels in basic HINet $s$ times.

\begin{table}
\centering
\tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|cc}
& \multicolumn{2}{|c}{SIDD~\cite{abdelhamed2018high}}\\
Method & PSNR & SSIM\\
\hline
DnCNN~\cite{zhang2017beyond}&23.66&0.583\\
MLP~\cite{burger2012image}&24.71&0.641\\
DM3D~\cite{dabov2007image}&25.65&0.685\\
CBDNet$^{\color{red}*}$~\cite{guo2019toward}&30.78&0.801\\
RIDNet$^{\color{red}*}$~\cite{anwar2019real}&38.71&0.951\\
AINDNet$^{\color{red}*}$~\cite{kim2020transfer}&38.95&0.952\\
VDN~\cite{yue2019variational}&39.28&0.956\\
SADNet$^{\color{red}*}$~\cite{chang2020spatial}&39.46&0.957\\
DANet+$^{\color{red}*}$~\cite{yue2020dual}&39.47&0.957\\
CycleISP$^{\color{red}*}$~\cite{Zamir2020CycleISP}&39.52&0.957\\
MPRNet~\cite{Zamir2021MPRNet}&39.71&0.958\\
\hline
HINet $0.5\times$(\textbf{ours})&\underline{39.82}&0.958\\
HINet (\textbf{ours})&\textbf{39.99}&0.958\\
\end{tabular}
\vspace{-0.2cm}
\caption{Denoising comparisons on SIDD~\cite{abdelhamed2018high} dataset. $^{\color{red}*}$ denotes the methods that use additional training data. Best and second best scores are \textbf{highlighted} and \underline{underlined}. Our HINet achieves 0.28 dB absolute improvement in PSNR over the previous best method MPRNet~\cite{Zamir2021MPRNet}.}
\label{tab.SIDD.results}
\vspace{-0.2cm}
\end{table}

\begin{table}
\centering
\tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|cc}
& \multicolumn{2}{|c}{GoPro~\cite{nah2017deep}}\\
Method & PSNR & SSIM\\
\hline
Xu et al.~\cite{xu2013unnatural} & 21.00 & 0.741\\
Hyun et al.~\cite{hyun2013dynamic} & 23.64 & 0.824\\
Whyte et al.~\cite{whyte2012non} & 24.60 & 0.846\\
Gong et al.~\cite{gong2017motion} & 26.40 & 0.863\\
DeblurGAN~\cite{kupyn2018deblurgan}& 28.70 & 0.858\\
Nah et al.~\cite{nah2017deep} & 29.08 & 0.914\\
Zhang et al.~\cite{zhang2018dynamic} & 29.19 & 0.931\\
DeblurGAN-v2~\cite{kupyn2019deblurgan} & 29.55 & 0.934\\
SRN~\cite{tao2018scale} & 30.26 & 0.934\\
Gao et al.~\cite{gao2019dynamic} & 30.90 & 0.935\\
DBGAN~\cite{zhang2020deblurring} & 31.10 & 0.942\\
MT-RNN~\cite{park2020multi} & 31.15 & 0.945\\
DMPHN~\cite{Zhang_2019_CVPR} & 31.20 & 0.940\\
Suin et al.~\cite{suin2020spatially} & 31.85 & 0.948\\
MPRNet~\cite{Zamir2021MPRNet} & \underline{32.66} & \textbf{0.959}\\
\hline
HINet (\textbf{ours}) & \textbf{32.77} & \textbf{0.959}\\
\end{tabular}
\vspace{-0.2cm}
\caption{Deblurring comparisons on GoPro~\cite{nah2017deep} dataset. Best and second best scores are \textbf{highlighted} and \underline{underlined}. Our HINet achieves 0.11 dB absolute improvement in PSNR over the previous best method MPRNet~\cite{Zamir2021MPRNet}.}
\label{tab.GoPro.results}
\vspace{-0.4cm}
\end{table}

\begin{table*}
\centering
\tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|cccccccccc|cc}
& \multicolumn{2}{c}{Test100~\cite{zhang2019image}}& \multicolumn{2}{c}{Rain100H~\cite{yang2017deep}}& \multicolumn{2}{c}{Rain100L~\cite{yang2017deep}}& \multicolumn{2}{c}{Test2800~\cite{fu2017removing}}& \multicolumn{2}{c}{Test1200~\cite{zhang2018density}}& \multicolumn{2}{|c}{Average}\\
Method&PSNR & SSIM&PSNR &SSIM &PSNR & SSIM & PSNR & SSIM & PSNR & SSIM & PSNR & SSIM\\
\hline
DerainNet~\cite{fu2017clearing}&22.77 & 0.810 & 14.92 & 0.592 & 27.03 & 0.884 & 24.31 & 0.861 & 23.38 & 0.835 & 22.48 & 0.796\\
SEMI~\cite{wei2019semi} & 22.35 & 0.788 & 16.56 & 0.486 & 25.03 & 0.842 & 24.43 & 0.782 & 26.05 & 0.822 & 22.88 & 0.744\\
DIDMDN~\cite{zhang2018density} & 22.56 & 0.818 & 17.35 & 0.524 & 25.23 & 0.741 & 28.13 & 0.867 & 29.65 & 0.901 & 24.58 & 0.770\\
UMRL~\cite{yasarla2019uncertainty} & 24.41 & 0.829 & 26.01 & 0.832 & 29.18 & 0.923 & 29.97 & 0.905 & 30.55 & 0.910 & 28.02 & 0.880\\
RESCAN~\cite{li2018recurrent}& 25.00 & 0.835 & 26.36 & 0.786 & 29.80 & 0.881 & 31.29 & 0.904 & 30.51 & 0.882 & 28.59 & 0.857\\
PreNe~\cite{ren2019progressive} & 24.81 & 0.851 & 26.77 & 0.858 & 32.44 & 0.950 & 31.75 & 0.916 & 31.36 & 0.911 & 29.42 & 0.897\\
MSPFN~\cite{jiang2020multi} & 27.50 & 0.876 & 28.66 & 0.860 & 32.40 & 0.933 & 32.82 & 0.930 & 32.39 & 0.916 & 30.75 & 0.903\\
MPRNet~\cite{Zamir2021MPRNet} & \underline{30.27} & \underline{0.897} & \underline{30.41} & \underline{0.890} & \underline{36.40} & \underline{0.965} & \underline{33.64} & \underline{0.938} & \underline{32.91} & \underline{0.916} & \underline{32.73} & \underline{0.921}\\
\hline
HINet (\textbf{ours})& \textbf{30.29}& \textbf{0.906} & \textbf{30.65} & \textbf{0.894} & \textbf{37.28} & \textbf{0.970} & \textbf{33.91} & \textbf{0.941} & \textbf{33.05} & \textbf{0.919} & \textbf{33.03} & \textbf{0.926}\\
\end{tabular}
\vspace{-.2cm}
\caption{Deraining comparisons on Test100~\cite{zhang2019image}, Rain100H~\cite{yang2017deep}, Rain100L~\cite{yang2017deep}, Test2800~\cite{fu2017removing} and Test1200~\cite{zhang2018density}. In addition, the average results over these datasets are provided. Best and second best scores are \textbf{highlighted} and \underline{underlined}. Our HINet achieves 0.3 dB absolute improvement in PSNR over the previous best method MPRNet~\cite{Zamir2021MPRNet}.}
\label{tab.Rain13k.results}
\vspace{-0.3cm}
\end{table*}

\begin{table}
\centering
\tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|l|c|rr|rr}
Dataset& Method & PSNR & \multicolumn{2}{c}{MACs(G)} &Time(ms)&speedup\\
\hline
\multirow{3}{4em}{SIDD~\cite{abdelhamed2018high}} & MPRNet~\cite{Zamir2021MPRNet} & 39.71 & 573.50&100\% & 78.8 & $1\times$\\
&HINet $0.5\times$&39.82&42.88&7.5\% & 11.6 & $6.8\times$\\
&HINet &39.99&170.71&29.8\% & 27.0& $2.9\times$\\
\hline
\multirow{2}{4em}{REDS~\cite{nah2019ntire}} & MPRNet~\cite{Zamir2021MPRNet} & 28.81 & 760.11&100\% & 90.1& $1\times$\\
&HINet  &28.83&170.71&22.5\% & 27.0 & $3.3\times$\\
\hline
\multirow{2}{4em}{GoPro~\cite{nah2017deep}} & MPRNet~\cite{Zamir2021MPRNet} & 32.66 & 760.11&100\% & 90.1 & $1\times$ \\
&HINet  &32.77&170.71&22.5\% & 27.0 & $3.3\times$\\
\hline
\multirow{2}{4em}{Rain13k} & MPRNet~\cite{Zamir2021MPRNet} & 32.73 &141.28&100\% & 37.4 & $1\times$\\
&HINet  &33.03&170.71&120.8\% & 27.0 & $1.4\times$\\
\end{tabular}
\vspace{-.2cm}
\caption{Comparing the PSNR and MACs of MPRNet~\cite{Zamir2021MPRNet} and ours. For Rain13k, we compare the average PSNR over Test100~\cite{zhang2019image}, Rain100H~\cite{yang2017deep}, Rain100L~\cite{yang2017deep}, Test2800~\cite{fu2017removing} and Test1200~\cite{zhang2018density}. MACs and Time are estimated with the input size of $1\times3\times256\times256$. The proportion of the calculations and the speedup compared to MPRNet~\cite{Zamir2021MPRNet} are also listed. Runtimes are computed with the Tesla V100 GPU.}
\label{tab.GmacsCompare}
\vspace{-.3cm}
\end{table}

\subsection{Main Results}
We show the effectiveness of HINet on different datasets in Table~\ref{tab.SIDD.results}, Table~\ref{tab.GoPro.results} and Table~\ref{tab.Rain13k.results}. In addition, we compare the MACs (\ie multiplier-accumulator operations) and inference time of MPRNet~\cite{Zamir2021MPRNet} and HINet in Table~\ref{tab.GmacsCompare}. MACs is estimated when the input is $1\times3\times256\times256$. Moreover, we conduct quality experiments to show the superiority of our method as shown in Figure~\ref{fig:visual_more}.

\paragraph{SIDD~\cite{abdelhamed2018high}} For image denoising, we train our model on the 320 high-resolution images of the SIDD dataset, and test on the 1280 patches from SIDD dataset. The results are shown in Table~\ref{tab.SIDD.results} and Table~\ref{tab.GmacsCompare}. Surprisingly, using only 7.5\% of MPRNet~\cite{Zamir2021MPRNet}'s MACs, our model exceeds it by 0.11 dB in PSNR. Moreover, our model exceeds MPRNet a big margin, 0.28 dB in PSNR under 30\% of the MACs, and is 2.9 times faster than MPRNet.

\paragraph{REDS~\cite{nah2019ntire} and GoPro~\cite{nah2017deep}} For image deblurring, we train our model on REDS~\cite{nah2019ntire} dataset with jpeg compression artifacts, and evaluate the results on REDS-val-300 as we described above. In addition, we conduct experiments on GoPro~\cite{nah2017deep} dataset for image deblurring following ~\cite{Zamir2021MPRNet,Zhang_2019_CVPR} \etc. It contains 2103 image pairs for training and 1111 pairs for evaluation. In Table~\ref{tab.GoPro.results} and Table~\ref{tab.GmacsCompare}, we compare our approach to the state-of-the-art methods. We get a comparable performance to MPRNet~\cite{Zamir2021MPRNet} with only 22.5\% MACs and $3.3\times$ speed advantage. It indicates the efficiency of our model.

\paragraph{Rain13k} For image deraining, we train our model on Rain13k as described above and evaluate results by Y channel in YCbCr color space following ~\cite{jiang2020multi,Zamir2021MPRNet}. It contains 13712 image pairs in the training set, and we evaluate the results on Test100~\cite{zhang2019image}, Rain100H~\cite{yang2017deep}, Rain100L~\cite{yang2017deep}, Test2800~\cite{fu2017removing}, Test1200~\cite{zhang2018density}. We show the results in Table~\ref{tab.Rain13k.results}. HINet achieves 0.3 dB absolute improvement in PSNR over the previous best method MPRNet~\cite{Zamir2021MPRNet} and $1.4\times$ faster than it.

\subsection{Ablation}The core idea of HINet lies in HIN Block. We evaluate it from multiple perspectives in this subsection.
It should be noted that these experiments are \emph{not} to achieve the performance of the state-of-the-art, but to illustrate the superiority of HIN on various models and image restoration tasks.
Therefore we mainly use the UNet of the first stage of HINet $0.5\times$, without skip connections between encoder and decoder for ablation experiments for fast feedback, and we denoted the UNet as ``HINet Simple'' in next. It is trained on $512\times512$ patches with a batch size of $32$ for $3\times10^5$ iterations. For optimizer, we follow the settings in ~\cite{Zamir2021MPRNet} except we set the learning rate to $4\times10^{-4}$ instead of $2\times10^{-4}$.

\paragraph{The effectiveness of Half Instance Normalization:}
We conduct experiments on various models and datasets, as we shown in Table~\ref{hin.effectiveness}.

On REDS~\cite{nah2019ntire} dataset (shown in Table~\ref{hin.effectiveness.REDS}), HIN brings 0.12 dB in PSNR for HINet Simple.
This illustrates the effectiveness of HIN on the REDS~\cite{nah2019ntire} dataset and HINet Simple model.

On GoPro~\cite{nah2017deep} dataset (shown in Table~\ref{hin.effectiveness.GoPro}), we reimplement the DMPHN(1-2-4-8)~\cite{Zhang_2019_CVPR}. We set the learning rate to $2\times10^{-4}$ and decreased to $1\times10^{-6}$ with cosine annealing strategy~\cite{loshchilov2016sgdr}. We train the model on $256\times256$ patches, with a batch size of 32 for $2\times10^5$ iterations. Flip and rotation are applied as data augmentation. HIN brings 0.42 dB boost in PSNR, it demonstrates HIN is effective on different models and different datasets.

We train PRMID~\cite{wang2020practical} and CycleISP~\cite{Zamir2020CycleISP} on SIDD~\cite{abdelhamed2018high} dataset as shown in Table~\ref{hin.effectiveness.SIDD}. For PRMID, we set the initial learning rate to $3\times10^{-3}$ with a batch size of 16 for $4\times10^5$ iterations. For CycleISP, we set the initial learning rate to $2\times10^{-4}$ with a batch size of 16 for $3\times10^5$ iterations. Adam optimizer, cosine annealing strategy~\cite{loshchilov2016sgdr} and flip/rotation as data augmentation are adopted in both cases. HIN brings 0.09 dB and 0.06 dB boost in PSNR on PRMID and CycleISP respectively. Since SIDD~\cite{abdelhamed2018high} and GoPro~\cite{nah2017deep}/REDS~\cite{nah2019ntire} are different image restoration task datasets, it demonstrates that HIN is effective in different image restoration tasks.

\begin{table}
    \tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{subtable}[h]{0.49\textwidth}
    \centering
      \begin{tabular}{lc|cc}
      &&\multicolumn{2}{c}{REDS~\cite{nah2019ntire}}\\
      Method&HIN?&PSNR&SSIM\\
      \hline
      HINet Simple&- &28.11&0.847\\
      HINet Simple&\checkmark &28.23&0.850\\

      \end{tabular}
      \vspace{.1cm}
      \caption{Comparison of the models with/without HIN on REDS~\cite{nah2019ntire} dataset for deblurring. HIN brings 0.12 dB in PSNR to HINet Simple. }
      \label{hin.effectiveness.REDS}
   \end{subtable}
   \begin{subtable}[h]{0.49\textwidth}
   \vspace{.2cm}
   \centering
   \tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{lc|cc}
      &&\multicolumn{2}{c}{GoPro~\cite{nah2017deep}}\\
      Method&HIN?&PSNR&SSIM\\
      \hline
      DMPHN(1-2-4-8)~\cite{Zhang_2019_CVPR}&- &30.98&0.943\\
      DMPHN(1-2-4-8)~\cite{Zhang_2019_CVPR}&\checkmark &31.40&0.948\\

      \end{tabular}
      \vspace{.1cm}
      \caption{Comparison of the models with/without HIN on GoPro~\cite{nah2017deep} dataset for deblurring. HIN brings 0.42 dB in PSNR to DMPHN(1-2-4-8)~\cite{Zhang_2019_CVPR}. It indicates HIN's effectiveness is robust to datasets and models.}
      \label{hin.effectiveness.GoPro}
   \end{subtable}
   \vspace{.15cm}
   \begin{subtable}[h]{0.49\textwidth}
   \centering
   \vspace{.25cm}
   \tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{lc|c}
   \vspace{.1cm}
   &&\multicolumn{1}{c}{SIDD~\cite{abdelhamed2018high}}\\
      Method&HIN?&PSNR\\
      \hline
      PRMID~\cite{wang2020practical}&- &39.30\\
      PRMID~\cite{wang2020practical}&\checkmark&39.39\\
      \hline
      CycleISP~\cite{Zamir2020CycleISP}&- &39.50\\
      CycleISP~\cite{Zamir2020CycleISP}&\checkmark&39.56\\
   \end{tabular}
   \vspace{.1cm}
   \caption{Comparison of the models with/without HIN on SIDD~\cite{abdelhamed2018high} dataset for denoising. HIN brings 0.09 dB and 0.06 dB in PSNR to PRMID~\cite{wang2020practical} and CycleISP~\cite{Zamir2020CycleISP} respectively. It indicates HIN's effectiveness is robust to image restoration tasks and model size.}
   \label{hin.effectiveness.SIDD}
   \end{subtable}
\vspace{-.1cm}
\caption{To demonstrate the effectiveness of HIN, we conduct experiments on various datasets and models.}
\vspace{-.2cm}
\label{hin.effectiveness}
\end{table}

\paragraph{Comparison with other Normalizations:}
Normalization has not been fully explored in image restoration tasks. We compare HIN with other normalizations to demonstrate the superiority of our proposed HIN. We conduct experiments on REDS~\cite{nah2019ntire} dataset on HINet Simple. The results are shown in Table~\ref{tab:vs.norm}. We denote the HINet Simple w/o. Norm as the baseline in the following. BN~\cite{ioffe2015batch} results a significant performance drop compares to baseline, \ie 0.12 dB in PSNR. We conjure that is because of the inaccurate batch statistics estimation when batch size is small. It is alleviated by SyncBN~\cite{zhang2018context} in some extent, \ie HINet Simple w/. SyncBN brings 0.1 dB gain compares to HINet Simple w/. BN. However, it is still inferior to baseline (28.09 dB vs. 28.11 dB in PSNR). As we can see, with the help of IN~\cite{ulyanov2017improved}, HINet Simple w/. IN exceed baseline by 0.03 dB in PSNR. It indicates IN facilitates the training of image restoration tasks. As shown in Table~\ref{tab:vs.norm}, HIN exceeds its counterparts, and it brings 0.12 dB gain in PSNR compares to baseline. It demonstrates the superiority of our proposed HIN.

\begin{table}[h]
    \centering
    \tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|c}
        Method &PSNR \\
        \hline
        HINet Simple w/o. Norm & 28.11 \\
        HINet Simple w/. BN~\cite{ioffe2015batch} & 27.99\\
        HINet Simple w/. SyncBN~\cite{zhang2018context} &28.09\\
        HINet Simple w/. LN~\cite{ba2016layer} & 28.09 \\
        HINet Simple w/. IN~\cite{ulyanov2017improved} & 28.14 \\
        HINet Simple w/. IBN~\cite{pan2018two} & 28.18 \\
        HINet Simple w/. GN~\cite{wu2018group} & 28.19\\
        \hline
        HINet Simple w/. HIN(\textbf{ours}) & \textbf{28.23}\\
    \end{tabular}

    \caption{Comparison of different normalization approaches on the image restoration task. Experiments are conducted on REDS~\cite{nah2019ntire} dataset. HINet Simple w/. BN means a fully normalized model with BN in the encoder.}

    \label{tab:vs.norm}
\end{table}

\paragraph{More training iterations:}
We analyze the impact of increasing the number of training iterations based on HINet Simple. We train the model on REDS~\cite{nah2019ntire} dataset for 300k, 600k, and 900k iterations with/without HIN respectively. The results are shown in Figure~\ref{fig:more_iters}. The gap between HINet Simple w/. HIN and HINet Simple w/o. HIN does not decrease as the number of iterations increased. We conjure that this is because HIN increases the upper limit of the model, not just speeds up the convergence.

\begin{figure}
    \centering
    \includegraphics[width=0.48\textwidth]{figures/more_iters.pdf}
    \caption{Effectiveness of HIN when training more iterations.}
    \label{fig:more_iters}

\end{figure}

\paragraph{Guideline of add HIN layer in an existing network:}
HINet Simple consists of 5 encoder blocks and 5 decoder blocks. We further explore the appropriate add location of HIN Block. The results are shown in Table~\ref{tab:HIN.position}. It indicates that adding HIN to all encoder blocks gets the highest score. In addition, adding HIN to the encoder and decoder causes performance drop \ie 28.23 dB to 28.21 dB. It demonstrates that adding more HIN does not necessarily lead to better performance. In practice, add one HIN layer to each encoder block might be a good choice.

\begin{table}[h]
    \centering
    \tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|ccccc|c|c}
    &\multicolumn{5}{c|}{Encoder}&\\
    Method&1&2&3&4&5&Decoder&PSNR\\

         \hline
         \multirow{12}{6em}{HINet Simple} &-&-&-&-&-&-&28.11\\
         &\checkmark&\checkmark&-&-&-&-&28.15\\
         &-&\checkmark&\checkmark&-&-&-&28.19\\
         &-&-&\checkmark&\checkmark&-&-&28.19\\
         &-&-&-&\checkmark&\checkmark&-&28.21\\
         &\checkmark&\checkmark&\checkmark&-&-&-&28.19\\
         &-&\checkmark&\checkmark&\checkmark&-&-&28.20\\
         &-&-&\checkmark&\checkmark&\checkmark&-&28.21\\
         &\checkmark&\checkmark&\checkmark&\checkmark&-&-&28.21\\
         &-&\checkmark&\checkmark&\checkmark&\checkmark&-&28.22\\
         &\checkmark&\checkmark&\checkmark&\checkmark&\checkmark&-&\textbf{28.23}\\
         &\checkmark&\checkmark&\checkmark&\checkmark&\checkmark&\checkmark&28.21\\
    \end{tabular}
    \caption{Guideline of add HIN layer in an existing network (\eg HINet Simple): adding HIN to all encoder blocks gets highest score, while more HIN does not necessarily lead to better performance.}

    \label{tab:HIN.position}
\end{table}

\subsection{Extension to HINet:}
In order to achieve better performance on NTIRE 2021 Image Deblurring Challenge Track2. JPEG artifacts~\cite{nah2021ntire}, we extend HINet, and adopt test time augmentation strategy. To further enhance the performance, we ensemble 3 similar models. In this subsection, we discuss the impact of these three methods on the results. And at the end, the results of the development phase and the test phase are provided. The results are evaluated on REDS-val-300, except the test phase result.

\paragraph{Wider, Deeper:}
It has been demonstrated that scaling up the model from width and depth improves the model capacity~\cite{tan2019efficientnet,yu2018wide}. For width, we simply use HINet $2\times$. For depth, we add two residual blocks at the end of each encoder block and decoder block.
It achieves a PSNR of 29.05 dB on REDS-val-300.

\paragraph{Test Time Augmentation and Ensemble:}
We adopt flip and rotation as test time augmentation. It brings about 0.14 dB in PSNR. In addition, we randomly crop hundreds of patches, randomly adopt flip and rotation augmentation on them. It brings about 0.05 dB in PSNR. We simply average the predictions of 3 models as a model ensemble. It brings about 0.01 dB in PSNR. With these strategies, our model boost PSNR from 29.05 dB to 29.25 dB.

\paragraph{Development phase result:}
For the development phase, we randomly crop 720 patches. The results are shown in Table~\ref{result.devphase}.

\begin{table}[]
    \centering
    \tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|ccc}
    Participants & PSNR & SSIM & rank\\
    \hline
        \textbf{ours} & \textbf{29.25} & \textbf{0.8190} & 1 \\
         participant A & \underline{29.17} & \underline{0.8183} & 2\\
         participant B & 29.14 & 0.8170 & 3\\
         participant C & 29.11 & 0.8171 & 4\\
         participant D & 29.10 & 0.8165 & 5\\
         participant E & 29.01 & 0.8141 & 6\\
         participant F & 28.94 & 0.8145 & 7\\
         participant G & 28.75 & 0.8093 & 8\\
         participant H & 28.68 & 0.8103 & 9\\
         participant I & 28.66 & 0.8082 & 10\\
    \end{tabular}
    \vspace{-0.1cm}
    \caption{Development phase result of NTIRE 2021 Image Deblurring Challenge Track 2. JPEG artifacts~\cite{nah2021ntire}. Best and second best scores are \textbf{highlighted} and \underline{underlined}. Our proposed method outperform others by 0.08 dB in PSNR.}
    \label{result.devphase}

\end{table}

\paragraph{Test phase result:}
For the test phase, we randomly crop 1000 patches. The results are shown in ~\ref{result.testphase}.

\begin{table}\centering
\tablestyle{5pt}{1.05}\setlength{\tabcolsep}{1.mm}\begin{tabular}{l|ccc}
Participants & PSNR & SSIM & rank\\
\hline
\textbf{ours} & \textbf{29.70} & \underline{0.8403} & 1\\
Noah\_CVLab & \underline{29.62} & 0.8397 & 2\\
CAPP\_OB & 29.60 & 0.8398 & 3\\
Baidu & 29.59 & 0.8381 & 4\\
SRC-B & 29.56 & 0.8385 & 5\\
Mier & 29.34 & 0.8355 & 6\\
VIDAR & 29.33 & \textbf{0.8565} & 7\\
DuLang & 29.17 & 0.8325 & 8\\
TeamInception & 29.11 & 0.8292 & 9\\
Giantpandacv & 29.07 & 0.8286 & 10\\
Maradona & 28.96 & 0.8264 & 11\\
LAB FHD & 28.92 & 0.8259 & 12\\
SYJ & 28.81 & 0.8222 & 13\\
Dseny & 28.26 & 0.8081 & 14\\
IPCV IITM & 27.91 & 0.8028 & 15\\
DMLAB & 27.84 & 0.8013 & 16\\
Blur Attack & 27.41 & 0.7887 & 17\\
\end{tabular}
\vspace{-0.1cm}
\caption{NTIRE 2021 Image Deblurring Challenge Track 2. JPEG artifacts result~\cite{nah2021ntire}. Best and second best scores are \textbf{highlighted} and \underline{underlined}. We exceed other participants over 0.08 dB in PSNR.}
\label{result.testphase}
\end{table}

\section{Conclusion}

In this work, we reuse Normalization in image restoration tasks. Specifically, we introduce Instance Normalization into a residual block and design an effective and efficient block: Half Instance Normalization Block (HIN Block). In HIN Block, we apply Instance Normalization for half of the intermediate features and keep the content information at the same time. Based on HIN Block, we further propose a multi-stage network called HINet. Between each stage, we use feature fusion and attention-guided map~\cite{zamir2021multi} across stages to ease the flow of information and enhance the multi-scale feature expression. Our proposed HINet surpasses the SOTA on various image restoration tasks. In addition, by using HINet, we won 1st place on the NTIRE 2021 Image De-blurring Challenge - Track2. JPEG Artifacts~\cite{nah2021ntire}.

\begin{figure*}[t]
    \centering
    \includegraphics[width=0.87\textwidth]{figures/visual_more_bak.png}
    \caption{More visualized results of HINet on various image restoration tasks. For each image pair, the left one is degraded and the right one is predicted by HINet.}
    \label{fig:visual_more}
\end{figure*}

\clearpage
{\small
\bibliographystyle{ieee_fullname}


\section{Introduction}

Optimal transport problem was initially formalized by the mathematician Gaspard Monge \cite{monge1781memoire}. Later a series of significant contribution in transportation theory leads to deep connections with more mathematical branches including partial differential equations, geometry, probability theory and statistics \cite{kantorovich1942translation,brenier1991polar}. Optimal transport provides a flexible framework for comparing probability measures. Monge and Kantorovich formulate the optimal transport problem in different ways, in which the Kantovorich formulation is a generalisation of Monge. For the Kantovorich's optimal transport problem, given two probability measures $\mu$ and $\nu$ defined on spaces $X$ and $Y$ respectively, and a cost function $c(x,y): X \times Y \to \mathbb{R}$, which measures the expense of moving one unit of mass from $x \in X$ to $y\in Y$, we aim at finding a joint distribution $\gamma^*$ defined on $Z = X \times Y$ such that the expectation of the cost over the joint distribution is minimized:
\begin{equation}\label{eq:ot}
  \inf_{\gamma} \cbr[2]{\int_{X\times Y}c(x,y) \,{\mathrm{d}}\gamma (x,y) \Big\vert ~ \gamma \in \Pi (\mu ,\nu ) } .
\end{equation}
The marginal constraints are given by
\begin{equation}
  \label{eq:Pi}
  \Pi(\mu, \nu) = \cbr[2]{\gamma\, \Big\vert \, \gamma\in\mathcal{P}(X\times Y ), \,  \int_{Y} \gamma \dif y = \mu,\, \int_{X} \gamma \dif x = \nu}.
\end{equation}
 Here $\mathcal{P}(X\times Y )$ denotes the set of probabilities defined on $X \times  Y$.  In this work, we only consider $X = Y = \mathbb{R}^d$. The optimal value of the objective in \eqref{eq:ot} is defined as the Wasserstein distance between probability measures $\mu$ and $\nu$.

In recent years, researchers in applied science fields also discover the importance of optimal transport. In spite of elegant theoretical results, generally computing Wasserstein distance is not an easy task. Computing the discrete optimal transport problem in a straightforward way leads to solving a linear programming problem whose computation cost can be unaffordable with large scale problem settings \cite{pele2009fast}.  In \cite{cuturi2013sinkhorn}, the author smooths the discrete OT problem with an entropic regularization, and designs a fast matrix scaling algorithm which demonstrates high efficiency. However, the computation can be even intractable when it comes to continuous case. In \cite{li2018parallel}, the authors compute the continuous problem by first discretizing the space, such treatment is unrealistic for many applications involving probabilistic measures lying in high dimensional space with complicated shapes.
We have witnessed the success of deep neural networks in dealing with the large scale continuous OT problem \cite{arjovsky2017wasserstein,seguy2017large}. But is it possible to save some efforts for parameter tuning and deal with the continuous problem from another perspective?

In this paper, instead of solving the standard continuous transport problem, we start with an entropy transport problem as a relaxed optimal transport problem with soft marginal constraints. Recently, the importance of entropy transport problem has drawn researchers' attention as people figuring out its duality connection with unbalanced optimal transport problem \cite{chizat2018unbalanced,liero2018optimal} and treat it as a canonical distance function on the space of positive measures \cite{liero2018optimal}. With these soft marginal constraints, we can realize the corresponding Wasserstein gradient flow as a time evolution Partial Differential Equation (PDE) and finally numerically solve the regularized problem by evolving an interacting particle system via Kernel Density Estimation techniques \cite{parzen1962estimation}.
To get samples from optimal coupling, the traditional methods like Linear Programming \cite{oberman2015efficient,schmitzer2016sparse,walsh2017general} or Sinkhorn \cite{cuturi2013sinkhorn} usually start with the discretization of the whole continuous space and compute the transport plan for discrete setting as the approximation of the continuous case. Our algorithm can directly output the sample approximation of the optimal coupling without any discretization or training process as neural network method \cite{seguy2017large,korotin2019wasserstein,makkuva2020optimal}. This is also very different from other traditional methods like Monge-Amp\`{e}re Equation \cite{benamou2014numerical} or dynamical scheme \cite{benamou2000computational,li2018parallel,ruthotto2020machine}. We note that a recent independent work \cite{daaloul2021sampling} on sampling algorithm for Wasserstein Barycenter problems shares similar ideas with our proposed method.

Our main contribution is to analyze the theoretical properties of the entropy transport problem constrained on probability space and construct the corresponding Wasserstein gradient flow. For the constrained transport problem, we prove the existence and uniqueness of the solution under certain assumptions, and further study the $\Gamma$-convergence property of the entropy transport functionals to the classical optimal transport functional. Based on the gradient flow we derive, we propose an innovative algorithm for obtaining the sample approximation for the optimal plan. Our method can deal with optimal transport problem between two known densities. As far as we know, despite the classical discretization methods \cite{benamou2000computational,benamou2014numerical,li2018parallel} there is no scalable way to solve this type of problem. We also demonstrate the efficiency of our method by numerical experiments.

The paper is structured as follows. We introduce the constrained entropy transport as a regularized optimal transport problem in section 2; We carry out the Wasserstein gradient flow approach and its particle formulation in section 3; We design the algorithm as an interacting particle system in section 4; and demonstrate its accuracy with empirical evidences in section 5 before the conclusion.

\section{Constrained Entropy Transport as a regularized Optimal Transport problem} \label{section introduction to ET}

\subsection{Entropy Transport problem}
In our research, we will mainly restrict our discussion on Euclidean Space (i.e. $\mathbb{R}^d$). We denote $\mathcal{M}(\mathbb{R}^d)$ as the space of finite positive Radon measures on $\mathbb{R}^d $. We denote $\mathcal{P}(\mathbb{R}^d)$ as the space of probability measures defined on $ \mathbb{R}^d$.

For $\mu,\nu\in\mathcal{M}(\mathbb{R}^d)$ and $\gamma\in\mathcal{M}(\mathbb{R}^d\times\mathbb{R}^d)$,
We denote $\pi_1:\mathbb{R}^d\times\mathbb{R}^d\rightarrow \mathbb{R}^d$ as the projection onto the first coordinate: $\pi_1(x,y)=x$; and $\pi_2$ as the projection onto the second coordinate. $\#$ represents the push-forward operation\footnote{Suppose $T:\mathbb{R}^m\rightarrow\mathbb{R}^n$ is a measurable map, suppose $p$ is a measure defined on $\mathbb{R}^m$. Then $T_{\#}p$ is a measure defined on $\mathbb{R}^n$ satisfying: $T_{\#}p(E)=p(T^{-1}(E))\quad\forall~\textrm{measurable set}~E\subset\mathbb{R}^n$.
}. Let us consider the following functional:
\begin{align}
 \mathcal{E}(\gamma|\mu,\nu) =  \iint_{\mathbb{R}^d\times\mathbb{R}^d}c(x,y)d\gamma(x,y) +D({\pi_1}_{\#}\gamma \|\mu)+D({\pi_2}_{\#}\gamma \|\nu).\label{general ET functional}
\end{align}
Here $c:\mathbb{R}^d\times\mathbb{R}^d\rightarrow [0,+\infty]$ is a lower semicontinuous cost function.
$D(\cdot\|\cdot):\mathcal{M}(\mathbb{R}^d)\times\mathcal{M}(\mathbb{R}^d)\rightarrow \mathbb{R}$ is the divergence functional defined as:
\begin{equation}
  D(\mu\|\nu) =
  \begin{cases}
  \int_{\mathbb{R}^d} F\left(\frac{d\mu}{d\nu}\right)d\nu \quad \textrm{if}~\mu\ll\nu\\
  +\infty \quad \textrm{otherwise}
  \end{cases}  \label{general divergence functional}
\end{equation}
Here $F:[0,+\infty)\rightarrow [0,+\infty]$ is some convex function and there exists at least one $s>0$ such that $F(s)<+\infty$.

Then the general \textbf{ Entropy Transport problem} can be formulated as:
\begin{equation}
 \inf\limits_{\gamma\in\mathcal{M}(\mathbb{R}^d\times\mathbb{R}^d)}\left\{ \mathcal{E}(\gamma|\mu,\nu) \right\}. \label{general ET}
\end{equation}
It is not hard to show that $\mathcal{E}$ is convex on $\mathcal{M}(\mathbb{R}^d\times\mathbb{R}^d)$:
\begin{theorem}\label{thm convexity of ET functional}
Under the previous assumptions on $c$ and $F$, for any$\gamma_a,\gamma_b\in\mathcal{M}(\mathbb{R}^d\times\mathbb{R}^d)$ and $0\leq t\leq 1$, we have:
\begin{equation*}
\mathcal{E}(t\gamma_a+(1-t)\gamma_b)\leq t\mathcal{E}(\gamma_a)+(1-t)\mathcal{E}(\gamma_b).
\end{equation*}
\end{theorem}
We have the following theorem on existence and uniqueness of minimizer for problem \eqref{general ET}:
\begin{theorem}\label{thm existence uniqueness sol }
  We consider problem \eqref{general ET} involving the entropy transport functional defined in \eqref{general ET functional}. Suppose that the cost $c$ and $F$ satisfy the previous assumptions. We further assume that there exists at least one $\gamma\in\mathcal{M}(\mathbb{R}^d\times\mathbb{R}^d)$ such that $\mathcal{E}(\gamma|\mu,\nu)<+\infty$. Then the problem \eqref{general ET} admits at least one optimal solution.\\
  If we further assume $c(x,y)=h(x-y)$ with strictly convex $h:\mathbb{R}^d\rightarrow[0,+\infty)$; $F$ is strictly convex, and is superlinear, i.e. $\lim_{s\rightarrow+\infty}\frac{F(s)}{s}=+\infty$; distribution $\mu$,$\nu$ has density functions, i.e. $\mu\ll\mathscr{L}^d,\nu\ll\mathscr{L}^d$ where $\mathscr{L}^d$ denotes the Lebesgue measure on $\mathbb{R}^d$. Under these further assumptions, there exists unique optimal solution to the problem \eqref{general ET}.
\end{theorem}

There could be many choices for cost $c$ and divergence $D$. For example, setting $c(x,y)=-\log\cos^2(\min    \{\frac{|x-y|}{2\tau},\pi  \})$ and $F(s)=s\log s-s+1$ leads to an entropy transport problem equivalent to solving for the Wasserstein-Fisher-Rao (or Hellinger–Kantorovich) distance between distributions $\mu$ and $\nu$ \cite{peyre2019computational}\cite{chizat2018unbalanced}\cite{liero2018optimal}.

In our research, we mainly treat the entropy transport problem \eqref{general ET} as a relaxed optimal transport problem with soft marginal constraints. Recall that optimal transport problem is formulated as:
\begin{equation}
  \inf_{\substack{\text{$\gamma\in\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)$, } \\ \text{ $\pi_{1 \#}\gamma=\mu,\pi_{2 \#}\gamma=\nu$} }}\iint c(x,y)~d\gamma(x,y) \label{OT}.
\end{equation}
Here $\mu,\nu\in\mathcal{P}(\mathbb{R}^d)$. \eqref{OT} can also be treated as an entropy transport problem:
\begin{align}
  &\inf_{\gamma\in\mathcal{M}(\mathbb{R}^d\times\mathbb{R}^d)}\left\{\mathcal{E}_{\iota}(\gamma|\mu,\nu)\right\} \label{iota ET}  \\ \textrm{with}\quad
  \mathcal{E}_{\iota}(\gamma|\mu,\nu) &= \iint c(x,y)d\gamma(x,y)
  +\int\iota\left(\frac{d\pi_{1 \#}\gamma}{d\mu}\right)d\mu + \int\iota\left(\frac{d\pi_{1 \#}\gamma}{d\nu}\right)d\nu\label{iota ET functional}
\end{align}
Here we choose $F(\cdot)=\iota(\cdot)$ in the original functional \eqref{general ET functional} with $\iota$ defined as:
\begin{equation*}
  \iota(s) = \begin{cases}
    0 \quad \textrm{if}~ s\neq 1\\
    +\infty \quad \textrm{if}~ s=1
  \end{cases}
\end{equation*}
It is not hard to verify that \eqref{iota ET} and \eqref{OT} are equivalent.

In order to derive a relaxed optimal transport problem as an entropy transport problem, we relax the divergence terms $D(\pi_{1 \#}\gamma\|\mu),D(\pi_{2 \#}\gamma \| \nu)$ involving marginal distributions of $\gamma$: For example, we may replace the function $\iota(\cdot)$ with $\Lambda F(\cdot)$, here $\Lambda>0$ is a large positive number and $F$ is a smooth convex function with $F(1)=0$ and 1 is the unique minimizer. There are definitely many choices for $F$, some popular choices are the power-like entropies \cite{liero2018optimal}:
\begin{equation*}
  F_\alpha(s) = \begin{cases}
   \frac{1}{\alpha(\alpha-1)}(s^\alpha-\alpha(s-1)-1) ~ \textrm{if}~ \alpha>1 \\
   s\log s -s +1 ~\textrm{if}~ \alpha = 1
  \end{cases}
\end{equation*}
In our research, we mainly focus on the case when $\alpha = 1$ since it is a canonical divergence functional in transportation theory and enables us to establish corresponding theoretical results. On the other hand, $\alpha = 1$ leads to more concise form when we are deriving for our algorithms and can reduce the computational cost. It also worth mentioning that when $\alpha=1$, the corresponding divergence $D(\cdot \| \cdot )$ is usually called Kullback-Leibler (KL) divergence \cite{kullback1951information} and we will denote it as $D_{\rm{KL}}(\cdot\|\cdot)$.

From now on, we should focus on the following functional involving cost $c$ and enforcing the marginal constraints by using KL-divergence:
\begin{equation}
  \mathcal{E}_{\Lambda,\rm{KL}}(\gamma|\mu,\nu) = \iint_{\mathbb{R}^d\times\mathbb{R}^d} c(x,y)~d\gamma(x,y) + \Lambda D_{\rm{KL}}(\pi_{1 \#} \gamma \| \mu ) + \Lambda D_{\rm{KL}}(\pi_{2 \#} \gamma \| \nu ). \label{KL ET functional}
\end{equation}
Since a majority of our applications are devoted to optimal transport problems with specific form of cost functions $c$ and marginals $\mu,\nu$, let us make the following assumptions in order to make our future discussion concise:
\begin{align}
   (A).~& c(x,y) = h(x-y)~\textrm{with}~ h ~\textrm{a strictly convex function}.\nonumber\\
   \quad (B). ~& \mu,\nu\in\mathcal{P}(\mathbb{R}^d)~\rm{and}~   \mu\ll\mathscr{L}^d, \nu\ll\mathscr{L}^d ~ \label{additional cond}
\end{align}
Here $\mathscr{L}^d$ is Lebesgue measure on $\mathbb{R}^d$.
\begin{theorem}\label{uniqueness optimal distribution OT}
 Under additional assumptions \eqref{additional cond}, the optimal transport problem \eqref{OT} admits a unique solution.
\end{theorem}
The proof of this theorem can be found in \cite{ambrosio2008gradient} Theorem 6.2.4.

\subsection{Constrained Entropy Transport problem and some of its properties} \label{subsec: PET}
We now restrict the functional $\mathcal{E}_{\Lambda,\rm{KL}}(\cdot|\mu,\nu)$ to the probability manifold $\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)$ instead of $\mathcal{M}(\mathbb{R}^d\times\mathbb{R}^d)$. There are mainly two reasons of such restriction:
\begin{itemize}
  \item This will allow us to compute the Wasserstein gradeint flow of $\mathcal{E}_{\Lambda,\rm{KL}}(\cdot|\mu,\nu)$ on probability manifold and will later lead to our algorithm in form of an interacting particle system (See section \ref{flow of ET});
  \item When restricting on $\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)$, there is a natural $\Gamma-$convergence of $\{\mathcal{E}_{\Lambda,\rm{KL}}( \cdot | \mu , \nu)\}_{\Lambda}$ to the entropy transport functional corresponding to Optimal Transport problem \eqref{OT} as $\Lambda\rightarrow +\infty$ (See Theorem \ref{thmGamma});
\end{itemize}
We thus consider the following optimization problem:
\begin{equation}
  \inf_{\gamma\in\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)}\left\{ \mathcal{E}_{\Lambda,\rm{KL}}(\gamma |\mu , \nu)  \right\}
  \label{constrained-ET}.
\end{equation}
In the following discussion, we will call such problem \eqref{constrained-ET} \textbf{constrained Entropy Transport problem} since we constrain the space for $\gamma$ on $\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)$.

We now denote $\mathcal{E}_{\min}=\inf_{\gamma\in\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)}\{\mathcal{E}_{\Lambda,\rm{KL}}(\gamma| \mu ,   \nu)\}$. It can be shown that this infimum value $\mathcal{E}_{\min}$ is finite, i.e. $\mathcal{E}_{\min}>-\infty$. The following theorem shows the existence of the optimal solution to problem \eqref{constrained-ET}. It also describes the relationship between the solution
to the constrained Entropy Transport problem \eqref{constrained-ET} and the solution to the general Entropy Transport problem \eqref{general ET}:
\begin{theorem}\label{main thm A}
  Suppose  $\tilde{\gamma}$ is the solution to original entropy transport problem \eqref{general ET}. Then we have $\tilde{\gamma}=Z\gamma$, here $Z = e^{-\frac{\mathcal{E}_{\rm{min}}}{2\Lambda}       } $ and $\gamma\in\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)$ is the solution to constrained Entropy Transport problem \eqref{constrained-ET}.
\end{theorem}
The following corollary
guarantees the uniqueness of optimal solution to \eqref{constrained-ET}:
\begin{corollary}\label{corollary uniqueness constrained - ET}
  The constrained Entropy Transport problem admits a unique optimal solution.
\end{corollary}
The proof of Theorem \ref{main thm A} and Corollary \ref{corollary uniqueness constrained - ET} are provided in Appendix \ref{App a2}.

We can now characterize the structure of the optimal solution to problem \eqref{constrained-ET}:
\begin{theorem}[Characterization of optimal distribution $\gamma_{cET}$ to problem \eqref{constrained-ET} ]\label{main part characterization of optimal gamma}
  We assume $\textrm{supp}(\mu)=\textrm{supp}(\nu)=\mathbb{R}^d$. Suppose $\gamma_{cET}$ solves the constrained Entropy Transport problem \eqref{constrained-ET}. Then there exist certain $\varphi,\psi \in B(\mathbb{R}^d;\mathbb{R})  $ satisfying:
  $\varphi(x)+\psi(y)\leq c(x,y)$ for any $x,y\in\mathbb{R}^d$, such that:
  \begin{align}
    \varphi(x)+\psi(y)=c(x,y) & \quad  \gamma_{cET}  - \textrm{almost surely}, \nonumber \\
    \pi_{1 \#}\gamma_{cET} =e^{\frac{   \mathcal{E}_{   \textrm{min } }  - 2\varphi }{2\Lambda}} \mu & \quad   \pi_{2 \#}\gamma_{cET}  = e^{\frac{   \mathcal{E}_{   \textrm{min } }  - 2\psi }{2\Lambda}}\nu \label{main part:marginal entropy transport on P}
  \end{align}
\end{theorem}
We provide a direct proof of this theorem in Appendix \ref{App a2}.
We can compare the structure of solution $\gamma_{cET}$ to \eqref{constrained-ET} with the solution $\gamma_{OT}$
to the Optimal Transport problem \eqref{OT}.
\begin{theorem}[Characterization of optimal distribution $\gamma_{OT}$ to problem \eqref{OT}]\label{main part characterization of optimal distribution to Optimal Transport}
If we assume additional condition on the cost function: $c(x,y)\leq a(x)+b(y)$ with $a\in L^1(\mu)$, $b\in L^1(\nu)$. Then there exists an optimal distribution $\gamma_{OT}$ to problem \eqref{OT}. There exist $\varphi,\psi\in C(\mathbb{R}^d)$ such that $\varphi(x)+\psi(y)\leq c(x,y)$ for any $x,y\in\mathbb{R}^d$ with:
\begin{align}
  \varphi(x)+\psi(y)=c(x,y)&\quad \gamma_{OT}-\textrm{almost surely}\nonumber  \\
  \pi_{1 \#}\gamma_{OT} = \mu& \quad \pi_{2 \#}\gamma_{OT} = \nu \label{main part:marginal optimal transport on P}
\end{align}
\end{theorem}
Since we are using constrained Entropy Transport problem \eqref{constrained-ET} to approximate Optimal Transport problem \eqref{OT}, we are interested in comparing the difference between their optimal distributions $\gamma_{c ET}$ and $\gamma_{OT}$. Although we can identify their difference from marginal conditions \eqref{main part:marginal entropy transport on P} and \eqref{main part:marginal optimal transport on P}, we currently do not have a quantitative analysis on the difference between $\gamma_{cET}$ and $\gamma_{OT}$. This may serve as one of our future research directions.

Despite the discussions on  the constrained problem \eqref{constrained-ET} with fixed $\Lambda$, we also establish asymptotic results for \eqref{constrained-ET} with quadratic cost $c(x,y)=|x-y|^2$ as $\Lambda \rightarrow +\infty$. For the rest of this section, we define:
\begin{equation*}
   \mathcal{P}_2(X) = \left\{\gamma \Big| ~\gamma\in\mathcal{P}(X),\gamma\ll\mathscr{L}^{ m},~\int_{X} |x|^2d\gamma<+\infty \right\}
\end{equation*}
where $X=\mathbb{R}^{m}$ is any $m$ dimensional Euclidean space.
Let us now consider $\mathcal{P}_2(\mathbb{R}^d\times\mathbb{R}^d)$ and assume it is equipped with the topology of weak convergence. We are able to establish the following $\Gamma$-convergence results for the functional $\mathcal{E}_{\Lambda,\rm{KL}}( \cdot|\mu,\nu)$ defined on $\mathcal{P}_2(\mathbb{R}^d\times\mathbb{R}^d)$:
\begin{theorem}[$\Gamma$-convergence]\label{thmGamma}
Suppose $c(x,y)=|x-y|^2$. Assume that we are given $\mu,\nu\in\mathcal{P}_2(\mathbb{R}^d)$ and at least one of $\mu$ and $\nu$ satisfies the Logarithmic Sobolev inequality with constant $ K >0$. Let  $\{\Lambda_n\}$ be a positive increasing sequence, satisfying $\lim_{n\rightarrow \infty}\Lambda_n = +\infty$. We consider the sequence of functionals $\{ \mathcal{E}_{\Lambda_n,\rm{KL}}(\cdot|\mu,\nu) \}$. Recall the functional $\mathcal{E}_{\iota}(\cdot|\mu,\nu)$ defined in \eqref{iota ET functional}. Then $\{ \mathcal{E}_{\Lambda_n,\rm{KL}}(\cdot|\mu,\nu)\}$ $\Gamma$- converges to $\mathcal{E}_{\iota}(\cdot|\mu,\nu)$ on $\mathcal{P}_2(\mathbb{R}^d\times\mathbb{R}^d)$.
\end{theorem}
We can further establish the equi-coercive property for the family of functionals $\{\mathcal{E}_{\Lambda_n,\rm{KL}}(\cdot|\mu,\nu)\}_n$ and we use the Fundamental Theorem of $\Gamma$-convergence \cite{dal2012introduction} \cite{braides2006handbook} to establish the following asymptotic results:
\begin{theorem}[Property of $\Gamma$-convergence]\label{main thm B}
Suppose $c(x,y)=|x-y|^2$. Assuming $\mu,\nu\in\mathcal{P}_2(\mathbb{R}^d)$ and both $\mu, \nu$ satisfy the Logarithmic Sobolev inequality with constants $ K_\mu,K_\nu>0$. According to Corollary \ref{thm existence uniqueness sol }, problem \eqref{constrained-ET} with functional $\mathcal{E}_{\Lambda_n,\rm{KL}}(\cdot | \mu,\nu   )$
admits a unique optimal solution, let us denote it as $\gamma_n$. According to Theorem \ref{uniqueness optimal distribution OT}, the Optimal Transport problem \eqref{OT} also admits a unique solution, we denote it as $\gamma_{OT}$. Then: $\lim_{n\rightarrow \infty}\gamma_n = \gamma_{OT}$ in $\mathcal{P}_2(\mathbb{R}^d\times \mathbb{R}^d)$.
\end{theorem}
Detailed proofs of these theorems regarding $\Gamma$-convergence are provided in Appendix\ref{App a3}.

\section{Wasserstein Gradient Flow Approach for Solving the Regularized Problem} \label{flow of ET}
\subsection{Wasserstein gradient flow}\label{Wass mfld and Wass grad flow}
 There are already numerous researches \cite{jordan1998variational,otto2001,ambrosio2008gradient} regarding Wasserstein gradient flows of different types of functionals defined on the Wasserstein manifold-like structure $(\mathcal{P}_2(\mathbb{R}^d),g^W)$ that successfully relate certain kinds of time evolution Partial Differential Equations (PDEs) to the manifold gradient of corresponding functionals on $(\mathcal{P}_2(\mathbb{R}^d), g^W   )$. The Wasserstein manifold-like structure is the manifold $\mathcal{P}_2(\mathbb{R}^d)$ equipped with a special metric $g^W$ induced by the 2-Wasserstein distance. Under this setting, the Wasserstein gradient flow of a certain functional $\mathcal{F}$ defined on $\mathcal{P}_2(\mathbb{R}^d)$ can thus be formulated as:
 \begin{equation}
    \frac{\partial\gamma_t}{\partial t} = -\textrm{grad}_W \mathcal{F}(\gamma_t)
    \label{Wasserstein Gradient Flow}
 \end{equation}
 One can explain this equation \eqref{Wasserstein Gradient Flow} as the continuous steepest descent algorithm applied to $\mathcal{F}$ in order to determine the minimizer of the target functional $\mathcal{F}$. For more detailed information regarding Wasserstein manifold-like structure and Wasserstein gradients, please check Appendix \ref{App b1}.

 \subsection{Wasserstein gradient flow of Entropy Transport functional}\label{sec3.2}
 We now come back to our constrained entropy transport problem \eqref{constrained-ET}. There are mainly two reasons why we choose to compute the Wasserstein gradient flow of functional $\mathcal{E}_{\Lambda,\rm{KL}}( \cdot |\mu,\nu )$:
 \begin{itemize}
   \item Computing the Wasserstein gradient flow is equivalent to applying gradient descent to determine the minimizer of the entropy transport functional \eqref{KL ET functional};
   \item In most of the cases, Wasserstein gradient flows can be realized as a time evolution PDE describing the density evolution of a stochastic process. As a result, once we derived the gradient flow, there will be a natural particle version associated to the gradient flow. And this will make the computation of gradient flow tractable since we can evolve the particle system by applying the forward Euler scheme.
 \end{itemize}
 Now let us compute the Wasserstein gradient flow of $\mathcal{E}_{\Lambda,\rm{KL}}(\cdot|\mu,\nu)$:
 \begin{equation}
    \frac{\partial\gamma_t }{\partial t}=-\textrm{grad}_W\mathcal{E}_{\Lambda,\rm{KL}}(\gamma_t|\mu,\nu) , ~~ ~   \gamma_t|_{t=0}=\gamma_0 \label{wass grad flow ET functional }
 \end{equation}
 To keep our notations concise, we denote $\rho(\cdot,\cdot,t) = \frac{d\gamma_t}{d\mathscr{L}^{2d}}$ $\varrho_1 = \frac{d\mu}{d\mathscr{L}^d}$, $\varrho_2 = \frac{d\nu}{d\mathscr{L}^d}$, we can show that the previous equation \eqref{wass grad flow ET functional } can be written as:
  \begin{equation}
   \frac{\partial \rho}{\partial t}=\nabla\cdot(\rho~\nabla (c(x,y)+ \Lambda \log(\frac{\rho_1(x,t)}{\varrho_1(x)}) + \Lambda \log(\frac{\rho_2(y,t)}{\varrho_2(y)}) ))
  \label{PDE vers}
 \end{equation}
 Here $\rho_1(\cdot,t) = \frac{d\pi_{1 \#}\gamma_t}{d\mathscr{L}^d} =\int\rho(\cdot,y,t)dy$ and $\rho_2(\cdot,t) = \frac{d\pi_{2   \#}\gamma_t}{d\mathscr{L}^d}=\int\rho(x,\cdot,t)dx$ are density functions of marginals of $\gamma_t$. We put the details of our derivation in
 Appendix \ref{App b2}.
 \begin{remark}\label{no displacement convexity}
   We are currently not clear about the displacement convexity of the functional $\mathcal{E}_{\Lambda, \rm{KL}  }(\cdot | \mu,\nu   )$ on Wasserstein manifold-like structure $(\mathcal{P}_2(\mathbb{R}^d\times\mathbb{R}^d),g^W)$, which will guarantee its gradient flow to converge at its minimizer. This will be one of our future research directions . In practice, we should rely on the computational results to tell us whether our method works properly.
 \end{remark}

 \subsection{Particle formulation of our derived gradient flow} \label{sectionSDE}
 Let us treat \eqref{PDE vers} as certain kind of continuity equation, i.e. we treat $\rho(\cdot,t)$ as the density of the time-evolving random particles. Then the vector field that drives the random particles at time $t$ should be $-\nabla(c(x,y)+\Lambda\log\left(\frac{\rho_1(x,t)}{\varrho_1(x)}\right)+\Lambda\log\left(\frac{\rho_2(y,t)}{\varrho_2(y)}\right))$. This helps us design the following dynamics $\{(X_t,Y_t)\}_{t\geq 0}$: (here $\dot X_t$ denotes the time derivative $\frac{d X_t}{dt}$)
 \begin{equation}
 \label{SDE version}
    \begin{cases}
    \dot X_t= -\nabla_x c(X_t,Y_t) + \Lambda (\nabla\log\varrho_1(X_t) -\nabla\log \rho_1(X_t,t)) ;\\
    \dot Y_t= -\nabla_y c(X_t,Y_t) + \Lambda (\nabla\log\varrho_2(Y_t) -\nabla\log \rho_2(Y_t,t)) ;
    \end{cases}
 \end{equation}
 where $\textrm{Law}(X_0,Y_0) =  \gamma_0 $. Here $\rho_1(\cdot,t)$ is the density of $\textrm{Law}(X_t)$ and $\rho_2(\cdot,t)$ is the density of $\textrm{Law}(Y_t)$.
 If we assume the process \eqref{SDE version} is well-defined, then the probability density $\rho_t(x,y)$ of $(X_t,Y_t)$ should solve the PDE \eqref{PDE vers}.

 When we take a closer look at \eqref{SDE version}, we can verify that the movement of particle $(X_t,Y_t)$ at certain time $t$ depends on the probability density of $\textrm{Law}((X_t,Y_t))$ at $(X_t,Y_t)$, which can be approximated by the distribution of the surrounding particles near $(X_t, Y_t)$.
 Such equation \eqref{PDE vers} can be treated as
 a limit case of aggregation-diffusion equation \cite{carrillo2019aggregation,carrillo2019blob} with Dirac kernel convolution.
Generally speaking, we plan to evolve \eqref{SDE version} as a particle aggregation model in order to converge to a sample-wised approximation of the Optimal Transport plan $\gamma_{OT}$ for OT problem \eqref{OT}.

 We expect that the distribution $\textrm{Law}((X_t,Y_t))$ converges to $\gamma_{ cET} $ as $t\rightarrow\infty$, here $\gamma_{ cET }  $ is the minimizer of functional $\mathcal{E}_{\Lambda,\rm{KL}}(\cdot|\mu,\nu)$ on $\mathcal{P}(\mathbb{R}^d\times\mathbb{R}^d)$. Recall the discussion in section \ref{subsec: PET}, $\gamma_{cET}   $ can be treated as an approximation to the solution of OT problem \eqref{OT}.

 In conclusion, evolving \eqref{SDE version} as an interacting particle system will provide a particle aggregation method for computing the sample-wise approximation to the Optimal Transport problem \eqref{OT}.

\section{Algorithmic Development}
\subsection{Numerical Approximation via Kernel Method}
To use the Euler scheme to simulate the stochastic process \eqref{SDE version}, we have to find a numerical approximation for the term $\nabla \log \rho(x)$. Here we use the Kernel Density Estimation \cite{parzen1962estimation} to approximate the density by convolving with kernel $K(x,\xi)$. In this paper, we simply choose the Radial Basis Function (RBF) kernel:
\begin{equation*}
  K(x,\xi)=\exp\left(-\frac{|x-\xi|^2}{2\tau^2}\right)
\end{equation*}
Then we approximate $\nabla\log\rho$ by:
\begin{equation}
\begin{split}
  \nabla\log\rho(x)\approx\nabla\log(K*\rho)(x)=\frac{(\nabla_x K)*\rho(x)}{K*\rho(x)} \label{kernel_approximation}
  \end{split}
\end{equation}
Here $K*\rho(x) = \int K(x,\xi)\rho(\xi)d\xi$, $(\nabla_xK)*\rho(x)=\int \nabla_x K(x,\xi)\rho(\xi)d\xi$\footnote{Notice that we always use $\nabla_x K$ to denote the partial derivative of $K$ with respect to the first components.}. Such technique is also known as blobing method, which was first studied in \cite{carrillo2019blob} and has already been applied to Bayesian sampling \cite{chen2018unified}. With this reformulation, we  can evaluate the gradient log density function based on the locations of the particles:
\begin{equation*}
    \frac{\mathbb{E}_{\xi\sim\rho}\nabla_x K(x,\mathbf{\xi})}{\mathbb{E}_{\xi\sim\rho}K(x,\xi)}\approx\frac{\sum_{k=1}^N\nabla_x K(x,\xi_k)}{\sum_{k=1}^N K(x,\xi_k)} ~~~\xi_1,...,\xi_N, \textrm{i.i.d.}\sim\rho
\end{equation*}

With the help of this method, we are able to construct the following particle system involving $N$ particles $\{(X_i , Y_i)\}_{i=1,...,N}$. For the $i$-th particle, we have:
\begin{equation}
\left\{
\begin{aligned}
\dot X_i(t) &= -\nabla_x c(X_i(t),Y_i(t))   - \Lambda \left(\nabla V_1(X_i(t)) +  \frac{\sum_{k=1}^N\nabla_x K(X_i(t),X_k(t))}{\sum_{k=1}^N K(X_i(t),X_k(t))} \right) \\
\dot Y_i(t) &= -\nabla_y c(X_i(t),Y_i(t))   - \Lambda \left(\nabla V_2(Y_i(t)) +  \frac{\sum_{k=1}^N\nabla_x K(Y_i(t),Y_k(t))}{\sum_{k=1}^N K(Y_i(t),Y_k(t))} \right)
\end{aligned}
\right. \label{num_particle_sys}
\end{equation}
Here we denote $V_1=-\log\varrho_1$, $V_2=-\log\varrho_2$. Since we only need the gradients of $V_1,V_2$ , our algorithm can deal with unnormalized continuous probability measures.
We numerically verify that when $t\rightarrow \infty$, the empirical distribution $\frac{1}{N}\sum_{i=1}^N\delta_{(X_i(t),Y_i(t))}$ will converge to the optimal distribution $\gamma_{ cET}$ that solves \eqref{constrained-ET} with sufficient large $N$ and $\Lambda$, while the rigorous proof is reserved for our future work.

\subsection{More Computational Efficiency with Random Batch Methods }
Taking closer look at the equation \eqref{num_particle_sys}, we can see the main computational efforts are put into approximating the gradient log density function. In each time step, the computational cost is of the order $\mathcal{O}(n^2)$. Inspired by \cite{jin2020random}, we apply the Random Batch Methods (RBM) here to reduce the computational cost. Assume that we have $n$ particles in the system and we can divide all particles into $m$ batches equally. Then in each iteration, we only consider the particles in the same batch as the particle $X_i$ when we evaluate the $\nabla \log \rho(X_i)$. From simple analysis we know that the computational cost now is reduced to the order $\mathcal{O}(n^2/m)$, which is a significant improvement. With proper choice of batch size, we can still get reasonable approximation but spend much less computational efforts. The algorithm scheme is summarized in the algorithm \ref{alg:alg_par_sys}.

\subsection{Extention to Wasserstein Barycenter Problem}
Our framework can be extend to multi-marginal problems. Suppose we have $m$ marginal distributions $\mu_1,...,\mu_m$
with cost function $c(x_1,...,x_m)$. The general multi-marginal problem \cite{gangbo1998optimal} can be formulated as:
\begin{equation}
 \inf_{\gamma\in\Pi(\mu_1,...,\mu_m)}\left\{\int_{\mathbb{R}^{m\times d}}c(x_1,...,x_m)~d\gamma(x_1,...,x_m)\right\}\label{multi marginal problem}
\end{equation}
Here $\Pi(\mu_1,...,\mu_m)$ denotes the set of $\gamma\in\mathcal{P}(\mathbb{R}^{m\times d})$ with its $m$ marginals equal to $\mu_1,...\mu_m$. To deal with \eqref{multi marginal problem}, we extend the entropy transport functional \eqref{KL ET functional} to:
\begin{align}
  \mathcal{E}(\gamma|\mu_1,...,\mu_m) \label{multi marginal functional}
  = \int_{\mathbb{R}^{m d}} c(x_1,...,x_m)
  d\gamma(x_1...x_m) + \sum_{j=1}^m \Lambda_j D_{\rm{KL}}(\gamma_j\|\mu_j)
\end{align}
Where $\gamma_j = \pi_{j\#}\gamma$. It is natural to extend the constrained Entropy Transport problem \eqref{constrained-ET} to the problem:
\begin{equation*}
  \min_{\gamma\in\mathcal{P}(\mathbb{R}^{m \times d})} \mathcal{E}(\gamma|\mu_1,...,\mu_m)
\end{equation*}
Similar to the two marginals case, we can derive the Wasserstein flow of functional \eqref{multi marginal functional} on $\mathcal{P}_2(\mathbb{R}^{m\times d})$ and compute its  corresponding particle flow in order to evaluate an approximation to the optimal solution of \eqref{multi marginal problem}.

We now consider applying our particle flow algorithm to Wasserstein barycenter problem \cite{agueh2011barycenters}\cite{cuturi2014fast}, which can be treated as a specific multi-marginal problem. This  barycenter problem is formulated as:
\begin{equation}
\min_{\mu\in\mathcal{P}_2(\mathbb{R}^d)} \sum_{i=1}^m \lambda_i W_2^2(\mu,\mu_i)\label{barycenter original}
\end{equation}
Here $\lambda_i>0$ are the weights. The barycenter problem has an equivalent multi-marginal formulation.
We consider the following multi-marginal problem:
\begin{equation}
  \min_{ \gamma\in\Pi(\mu_1,...,\mu_m)
  } \left\{\int_{\mathbb{R}^{(m+1)\times d}} \sum_{i=1}^m \lambda_i|x-x_i|^2~d\gamma(x,x_1,...,x_m) \right\} \label{barycenter multimargin}
\end{equation}
Notice that there is no marginal constraint for $\gamma$ on the 0-th "$x$" component. \eqref{barycenter original} and \eqref{barycenter multimargin} are equivalent\cite{agueh2011barycenters} in the following sense : if  $\bar{\mu}$ and $\bar{\gamma}$ are the optimal solutions to \eqref{barycenter original} and \eqref{barycenter multimargin}, then $\pi_{0\#}\bar{\gamma}=\bar{\mu}$.

We now apply our scheme to solve \eqref{barycenter multimargin}. We need to consider the functional:
\begin{align*}
  \mathcal{E}(\gamma|\mu_1,..,\mu_m)=\int_{\mathbb{R}^{(m+1)\times d}} \sum_{j=1}^{m} \lambda_j|x-x_j|^2 ~ d\gamma(x,x_1,...,x_m) +\sum_{j=1}^m \Lambda_j D_{\rm{KL}}(\gamma_j\|\mu_j)
\end{align*}
where $\gamma_j = \pi_{j \#}\gamma$, $j=1,...,m$. The particle system of the Wasserstein gradient flow of this functional can be written as:
\begin{align*}
\begin{cases}
\dot X^{(0)}_i(t) &= \sum_{j=1}^{m} -2\lambda_j(X_i^{(0)}(t)-X_i^{(j)}(t))\\

\dot X^{(1)}_i(t) &= -\nabla_{x_1} d(X^{(0)}_i(t), X^{(1)}_i(t))
\quad + \Lambda_1 \left( \nabla_x \log\varrho_1(X^{(1)}_i(t)) \right. \\
&\quad\left.+ \frac{\sum_{k=1}^N\nabla_\xi K(X^{(1)}_i(t),X^{(1)}_k(t))}{\sum_{k=1}^N K(X^{(1)}_i(t),X^{(1)}_k(t))}\right) \\
&...\\
\dot X^{(m)}_i(t) &= -\nabla_{x_m} d(X^{(0)}_i(t),X^{(m)}_i(t))
\quad + \Lambda_m \left(\nabla_x \log\varrho_m(X^{(m)}_i(t))\right.\\
&\quad\left.+\frac{\sum_{k=1}^N\nabla_\xi K(X^{(m)}_i(t),X^{(m)}_k(t))}{\sum_{k=1}^N K(X^{(m)}_i(t),X^{(m)}_k(t))}\right)
\end{cases}
\label{bary center dynamical_system_marginal_density}
\end{align*}
here $i$ goes from $1$ to $N$ . There are $N$ particles $\{(X_i^{(0)},X_i^{(1)},...,X_i^{(m)})\}_{i=1,...,N}$ evolving in $\mathbb{R}^{(m+1)\times d}$ together. As $t\rightarrow \infty$, the empirical distribution of particles $\{X^{(0)}_i\}_{i=1,...,N}$ are expected to be an approximation of barycenter $\bar{\mu}$ of distributions $\mu_1,...,\mu_m$.

\begin{algorithm}[tb]
\caption{Random Batch Particle Evolution Algorithm}
\label{alg:alg_par_sys}
\begin{algorithmic}
   \STATE {\bfseries Input:} The density functions of the marginals $\varrho_1, \varrho_2$, timestep $\Delta t$, total number of iterations $T$, parameters of the chosen kernel $K$
   \STATE {\bfseries Initialize:} The initial locations of all particles $X_i(0)$ and $Y_i(0)$ where $i=1,2,\cdots,n$,
   \FOR{t\,=\,1,2,$\cdots$,T}
   \STATE Shuffle the particles and divide them into $m$ batches: $\mathcal{C}_1,\cdots,\mathcal{C}_m$
    \FOR{each batch $\mathcal{C}_q$}
    \STATE Update the location of each particle $(X_i,Y_i)$ ($i\in \mathcal{C}_q$)
    \begin{align*}

    \boldsymbol{X}_i &\leftarrow\boldsymbol{X}_i - \Delta t[~ \nabla_x c(X_i(t),Y_i(t)) + \Lambda \nabla V_1(X_i(t))
              +\Lambda \frac{\sum_{k\in \mathcal{C}_q}\nabla_x K(X_i(t),X_k(t))}{\sum_{k\in \mathcal{C}_q} K(X_i(t),X_k(t))} ~] \\
    \boldsymbol{Y}_i &\leftarrow\boldsymbol{Y}_i - \Delta t[~ \nabla_y c(X_i(t),Y_i(t)) + \Lambda \nabla V_2(Y_i(t))
              +\Lambda \frac{\sum_{k\in \mathcal{C}_q}\nabla_x K(Y_i(t),Y_k(t))}{\sum_{k\in \mathcal{C}_q} K(Y_i(t),Y_k(t))} ~]
    \end{align*}
    \ENDFOR
   \ENDFOR
   \STATE {\bfseries Output:} A sample approximation of the optimal coupling: $X_i(T), Y_i(T)$ for $i=1,2,\cdots,n$
\end{algorithmic}
\end{algorithm}

\section{Numerical Experiments}
In this section, we test our algorithm on several toy examples. All experiments are conducted on a machine with 2.20GHz CPU, 16GB of memory.

{\bf 1D Gaussian} We set two 1D Gaussian distributions $\mathcal{N}(x;-4,1), \mathcal{N}(x;6,1)$ as marginals and run the algorithm to compute the sample approximation of the optimal transport plan between them. We set $\lambda=200, \Delta t=0.001$ and run it with 1000 particles $(X_i,Y_i)$'s for 1000 iterations. We initialize the particles by drawing 1000 i.i.d. sample points from $\mathcal{N}(x;-20,4)$ as $X_i$'s and 1000 i.i.d. sample points from $\mathcal{N}(x;20,2)$ as $Y_i$'s. The empirical results are shown in figure \ref{fig:1dGauss_marg} and figure \ref{fig:1dGauss_map}. We can see that after 1000 iterations, we get a good sample approximation of the optimal transport plan.

\begin{figure}
\centering
{\scriptsize
\begin{tabular}{ccc}
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_marg_iter_0.pdf} &
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_marg_iter_5.pdf} &
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_marg_iter_10.pdf}\\
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_marg_iter_15.pdf} &
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_marg_iter_20.pdf} &
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_marg_iter_25.pdf}\\
\end{tabular}
}
\caption{Marginal plot for 1D Gaussian example. The red and black dashed lines correspond to two marginal distribution respectively and the solid blue and green lines are the kernel estimated density functions of particles at certain iterations. After first 25 iterations, the particles have matched the marginal distributions very well.}
\label{fig:1dGauss_marg}
\end{figure}

\begin{figure}
\centering
{\scriptsize
\begin{tabular}{ccc}
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_map_iter_0.pdf} &
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_map_iter_500.pdf} &
\includegraphics[width=.28\textwidth]{1dGauss/1dGauss_map_iter_1000.pdf} \\
\end{tabular}
}
\caption{The sample approximation for 1D Gaussian example. The orange dash line corresponds to the optimal transport map $T(x)=x+10$.}
\label{fig:1dGauss_map}
\end{figure}

{\bf 1D Gaussian Mixture}
Then we apply the algorithm to two 1D Gaussian mixture $\varrho_1=\frac{1}{2}\mathcal{N}(x;-1,1)+\frac{1}{2}\mathcal{N}(x;1,1), \varrho_2=\frac{1}{2}\mathcal{N}(x;-2,1)+\frac{1}{2}\mathcal{N}(x;2,1)$. For experiment, we set $\lambda=60, \Delta t=0.0004$ and run it with 1000 particles $(X_i,Y_i)$'s for 5000 iterations. We initialize the particles by drawing 2000 i.i.d. sample points from $\mathcal{N}(x;0,2)$ as $X_i$'s and $Y_i$'s. In figure \ref{fig:1dGauss_mix}, we can see the particles still match the marginal distributions well and give a clear approximation for the optimal transport map.

\begin{figure}
\centering
{\scriptsize
\begin{tabular}{cc}
\includegraphics[width=.42\textwidth]{1dGauss_mix/1dGauss_mix_marg.pdf} &
\includegraphics[width=.42\textwidth]{1dGauss_mix/1dGauss_mix_map.pdf} \\
\end{tabular}
}
\vspace{-6pt}
\caption{1D Gaussian mixture. Left. Marginal plot. The dash lines correspond to two marginal distributions. The histogram indicates the distribution of particles after 5000 iterations. Right. Sample approximation for the optimal coupling.}
\label{fig:1dGauss_mix}
\end{figure}

{\bf Synthetic 2D Data} Given two marginals $\mu, \nu$ and cost function $|x-y|^p$, we can get a constant speed geodesic connecting two marginals by defining the curve $\mu_t = (\pi_t)_{\#} \gamma$ where $\gamma$ is the optimal coupling and $\pi_t(x,y) = (1-t)x+t y$. Given two gray scale images, if we normalize the pixel intensity, the image can be treated as a histogram representing a discrete 2D distribution. By applying the RBF kernel, the image can be converted to a continuous distribution as Gaussian mixture. Since our method gives a sample approximation for the optimal coupling, we are able to get the sample approximation of a series of distributions interpolating between two given marginals. In figure \ref{fig:marg}, we plot several simple gray scale images which are converted to continuous probability densities and used as marginals in our experiments. In figures \ref{fig:star_circle_mnist}
, we show two examples of transporting one point cloud image to the other.
\vspace{-6pt}
\begin{figure}[htp]
\centering
{\scriptsize
\begin{tabular}{cccc}
\includegraphics[width=.14\textwidth]{syn1/syn1_marg1.png} &
\includegraphics[width=.14\textwidth]{syn1/syn1_marg2.png} &
\includegraphics[width=.14\textwidth]{syn2/syn2_marg1.png} &
\includegraphics[width=.14\textwidth]{syn2/syn2_marg2.png} \\
(a) & (b) & (c) & (d)
\end{tabular}
}
\caption{Some gray scale images.}
\label{fig:marg}
\end{figure}

\begin{figure}[htp]
\centering
{\scriptsize
\begin{tabular}{cccccc}
\includegraphics[width=.14\textwidth]{syn1/syn1_1} &
\includegraphics[width=.14\textwidth]{syn1/syn1_2} &
\includegraphics[width=.14\textwidth]{syn1/syn1_3} &
\includegraphics[width=.14\textwidth]{syn1/syn1_4} &
\includegraphics[width=.14\textwidth]{syn1/syn1_5} &
\includegraphics[width=.14\textwidth]{syn1/syn1_6} \\
\includegraphics[width=.14\textwidth]{syn2/syn2_1} &
\includegraphics[width=.14\textwidth]{syn2/syn2_2} &
\includegraphics[width=.14\textwidth]{syn2/syn2_3} &
\includegraphics[width=.14\textwidth]{syn2/syn2_4} &
\includegraphics[width=.14\textwidth]{syn2/syn2_5} &
\includegraphics[width=.14\textwidth]{syn2/syn2_6}
\end{tabular}
}
\caption{The sample approximation of a sequence of distributions interpolating between the two given distributions: "star" and "disks", MNIST handwritten digits "6" and "1".}
\label{fig:star_circle_mnist}
\end{figure}

{\bf Wasserstein Barycenters}
As we discuss in the previous section, we can numerically solve the Wasserstein barycenter problem using our scheme. Given two Gaussian distributions $\rho_1 = \mathcal{N}(-10,1), \rho_2 = \mathcal{N}(10,1)$, and cost function
\begin{equation*}
    c(x,x_1,x_2) = w_1 \|x-x_1\|^2 + w_2 \|x-x_2\|^2,
\end{equation*} we can compute sample approximation of the barycenter $\bar{\rho}$ of $\rho_1,\rho_2$. We try different weights $[w_1,w_2]=[0.25,0.75],[0.5,0.5],[0.75,0.25]$ to test our algorithm. The experimental results are shown in fig \ref{fig:baryc}. The distribution of the particles corresponding to the barycenter random variable $X_0$ converges to $\mathcal{N}(5,1), \mathcal{N}(0,1), \mathcal{N}(-5,1)$ successfully after 2000 iterations, which demonstrates the accuracy of the algorithm.

\begin{figure}[htp]
\centering
{\scriptsize
\begin{tabular}{ccc}
\includegraphics[width=.28\textwidth]{barycenter/barycenter_13.pdf} &
\includegraphics[width=.28\textwidth]{barycenter/barycenter_55.pdf} &
\includegraphics[width=.28\textwidth]{barycenter/barycenter_31.pdf}
\end{tabular}
}
\caption{Density plots for 1D Wasserstein barycenter example.  The red dashed lines correspond to two marginal distributions respectively and the solid green lines are the kernel estimated density functions of the particles $X_1$'s and $X_2$'s. The solid blue line represents the kernel estimated density function of the particles corresponding to the barycenter. Left. $[w_1,w_2]=[0.25,0.75]$. Middle. $[w_1,w_2]=[0.5,0.5]$. Right. $[w_1,w_2]=[0.75,0.25]$.}
\label{fig:baryc}
\end{figure}

\section{Conclusion}
We propose the constrained Entropy Transport problem \eqref{constrained-ET} and study its theoretical properties. We discover that the optimal distribution of \eqref{constrained-ET} can be treated as an approximation to the optimal plan of the Optimal Transport problem \eqref{OT} in the sense of $\Gamma$-convergence . We also construct the Wasserstein gradient flow of the Entropy Transport functional. Based on that, we propose an innovative algorithm which iteratively evolves a particle system to compute for the sample-wised optimal distribution to the constrained Entropy Transport problem \eqref{constrained-ET}.

For future work, on theoretical aspect, we will mainly concentrate on the quantitative study of the discrepancy between $\gamma_{cET}$ and $\gamma_{OT}$ and the analysis of displacement convexity of functional $\mathcal{E}_{\Lambda,\rm{KL}}(\cdot|\mu,\nu)$. On numerical aspect, we will focus more on producing further examples in higher dimensional space and finding potential applications of our method to different areas of machine learning research.
\bibliographystyle{splncs04}