# Progressive Energy-Based Cooperative Learning for Multi-Domain Image-to-
Image Translation

Weinan Song  University of California, Los Angeles  Yaxuan Zhu  University of
California, Los Angeles  Lei He  Eastern Institute for Advanced Study, Ningbo,
China  University of California, Los Angeles  Yingnian Wu  University of
California, Los Angeles  Jianwen Xie  Akool Research

###### Abstract

This paper studies a novel energy-based cooperative learning framework for
multi-domain image-to-image translation. The framework consists of four
components: descriptor, translator, style encoder, and style generator. The
descriptor is a multi-head energy-based model that represents a multi-domain
image distribution. The components of translator, style encoder, and style
generator constitute a diversified image generator. Specifically, given an
input image from a source domain, the translator turns it into a stylised
output image of the target domain according to a style code, which can be
inferred by the style encoder from a reference image or produced by the style
generator from a random noise. Since the style generator is represented as an
domain-specific distribution of style codes, the translator can provide a one-
to-many transformation (i.e., diversified generation) between source domain
and target domain. To train our framework, we propose a likelihood-based
multi-domain cooperative learning algorithm to jointly train the multi-domain
descriptor and the diversified image generator (including translator, style
encoder, and style generator modules) via multi-domain MCMC teaching, in which
the descriptor guides the diversified image generator to shift its probability
density toward the data distribution, while the diversified image generator
uses its randomly translated images to initialize the descriptor’s Langevin
dynamics process for efficient sampling. We also bring in two regularization
strategies for both the descriptor and the translator to significantly improve
the cooperative learning. To further enhance the efficiency and scalability,
we propose a progressive cooperative learning strategy to train our framework
from low resolution to high resolution. Strong empirical results on CelebA-HA
and AFHQ datasets are shown to verify the effectiveness of our energy-based
image translation framework.

##  1 Introduction

The task of image-to-image translation primarily involves the learning of
mappings between different visual domains. This learning framework carries
immense application value in the field of generative artificial intelligence,
facilitating the development of various creative products for artificial
intelligence-generated contents (AIGC). In this context, a “domain” refers to
a collection of images belonging to a visually distinctive category such as
the gender of a person and animal species. Within each domain, every image
exhibits a unique appearance, encompassing image-specific elements such as
hairstyle and makeup, commonly referred to as its “style”. An ideal image-to-
image translation framework should possess the ability to handle multiple
domains, efficiently process high-resolution images, and provide diverse
synthesis (i.e., one-to-many mapping) when translating to each target domain.
While several pioneering work [1, 2] have made significant strides towards
this goal, it is noteworthy that they predominantly rely on adversarial
learning techniques. In contrast, our paper explores the multi-domain image-
to-image translation problem using energy-based modeling techniques. This
research is motivated by the current swift advancements in energy-based
modeling and learning technologies.

By leveraging the representation power of an energy-based model and the
sampling efficiency of a latent variable model, the Generative Cooperative
Network [3], also known as CoopNets, and its variants [4, 5], have achieved
impressive results in numerous computer vision tasks, such as image generation
[3, 4, 5], visual salient object detection[6], supervised image-to-image
translation [7], and unsupervised image-to-image translation [8]. However,
while the cross-domain translation framework, CycleCoopNets [8], has
demonstrated success in unpaired image-to-image translation, it is only
capable of learning the relation between two different domains at a time. Such
an approach has a limited scalability in deal with multiple domains, as a
separate model must be trained for each pair of domains. In addition,
cooperative learning still faces challenges when it comes to translating high-
resolution images. This is because the translation process involves sampling
from the energy-based model via Langevin dynamics, which can be difficult to
apply to high-resolution image spaces. To tackle the aforementioned challenges
existing in the current cooperative learning (or more generally, energy-based
learning) for multi-domain unsupervised image-to-image translation, this paper
proposes a novel cooperative learning framework, MD-CoopNets, to ensure
scalability, flexibility, stability and efficiency for applying energy-based
framework to image-to-image translation.

To be specific, the MD-CoopNets consists of four components: descriptor,
translator, style generator and style encoder. (1) The descriptor is a multi-
head energy-based model that represents a multi-domain image distribution,
where each head of the energy function corresponds to one image domain. (2)
The style generator is a multi-head latent variable model responsible for
generating domain-specific style codes. It achieves this by transforming a
Gaussian latent code into a style code. Each head in the style generator
corresponds to one specific domain. It represents a domain-specific
distribution of style codes. (3) The style encoder extracts domain-specific
style codes from an input image using a multi-head encoder. Each head of the
encoder corresponds to a specific domain. (4) The translator is a style-
controlled mapping, which takes an image and a style code as input, and then
transforms the image into a translated image that reflects the desired style
indicated by the style code. The style code can be obtained either from the
style generator or the style extractor. The style generator, style encoder,
and style-controlled translator can constitute a diversified image generator.

As to the learning, the multi-domain descriptor and the diversified image
generator engage in a cooperative game, where the multi-domain descriptor
guides the diversified image generator in aligning its mapping towards the
target domains using MCMC teaching, while the image generator assists in
expediting the descriptor’s MCMC teaching process by providing a good
initialization. Specifically, to enforce a meaningful latent space of style
codes, we train the style-controlled translator and style encoder by
reconstructing style codes that are randomly generated from the style
generator. To enforce translated image to preserve the domain-invariant
property of the input reference image, we train the translator with a cycle
consistency loss. To enforce the one-to-many translation output, we regularize
the translator via a diversity sensitive loss, such that, given an identical
reference image, different style codes can lead to sufficiently diversified
translated outputs.

Additionally, we propose to improve the cooperative learning algorithm by
incorporating some loss terms to regularize the behaviors of the components in
our framework. Firstly, we put an $l_{2}$ regularization on the output of the
energy function of the descriptor to limit the magnitude of the energy values.
To accelerate and stabilize the teaching process provided by the descriptor’s
MCMC, we propose to use the energy function to regularize the output of the
translator. These regularization techniques significantly improve the
performance of the cooperative learning.

To enhance efficiency, stability, and scalability, we present a novel
progressive cooperative learning algorithm for our model. Our approach
involves a gradual expansion of the components, initially operating on simpler
low-resolution images. As the cooperative training proceeds, new layers are
added to each component, enabling the model to handle more challenging high-
resolution images. This progressive growth strategy significantly accelerates
and stabilizes both training and sampling processes at higher resolutions.
Moreover, it offers the flexibility and convenience to scale up the resolution
of any pre-trained MD-CoopNets.

We demonstrate the effectiveness of our proposed multi-domain translation
model on the CelebA-HQ [9] and AFHQ [2] dataset. The translated examples
exhibit high fidelity and are comparable to GAN-based multi-domain translation
models. Furthermore, our progressive learning strategy improves the efficiency
and stability of the original training process, particularly when it comes to
translating high-resolution images.

Our contributions are listed as follows:

  * •

We propose a novel energy-based cooperative learning framework for multi-
domain image-to-image translation. We build a single multi-head energy-based
model to represent probability distributions of multiple domains, and train it
with a translator, a style encoder, and a style generator using a cooperative
manner.

  * •

We present a novel progressive learning algorithm to optimize the training
efficiency of our framework. Our approach adopts a progressive growth
strategy, advancing all components from low resolution to high resolution. It
yields a significant reduction in the total number of MCMC steps required for
training and sampling from the high-resolution model.

  * •

We propose regularization strategies to stabilize the cooperative learning,
which include an energy-based regularization loss for the translator and a
$l_{2}$ regularization loss for limiting the magnitude of the energy values of
the descriptor. Significant performance gain are obtained from these
regularization.

  * •

We demonstrate strong empirical results on CelebA-HQ and AFHQ datasets to
verify the proposed energy-based framework. Our method obtains state-of-the-
art performance among existing energy-based image translation models and
comparable results with GAN-based baselines.

##  2 Related Work

#### Energy-based Learning

Training energy-based models (EBMs) [10, 11, 12] involves maximizing the
likelihood of the observed data by adjusting the model’s energy function
parameters, which typically requires Markov chain Monte Carlo (MCMC) sampling
to evaluate the intractable gradient [13, 14, 15]. Contrastive divergence (CD)
[16, 17] is an efficient approximation algorithm for training energy-based
models by initializing the MCMC chains with observed data. [14] uses a noise-
initialized non-convergent short-run MCMC to train an EBM, and obtains a valid
flow-like generator trained with moment matching estimation. [18] defines a
sequence of conditional EBMs and forms a denoising diffusion process. To avoid
MCMC, [19] brings in normalizing flow and trains an EBM by flow contrastive
estimation. Learning an amortized sampler [20, 3, 21, 22, 23, 24] for EBMs is
also an alternative strategy. Our method has a single multi-head EBM to
represent multi-domain data distribution, and the image-to-image translator
serves as a multi-domain amortized sampler for the EBMs.

#### Cooperative Learning

Cooperative learning for energy-based models with MCMC teaching is first
proposed in [3], where the authors utilize an energy-based model as the
descriptor and a latent variable model as the generator to speed up the
learning of each other by maximum likelihood algorithms. During each training
iteration, the descriptor generates samples by finite-step MCMC sampling with
initialization by the generation from the generator for maximum likelihood
estimation. Simultaneously, the sampling results from descriptor are used to
directly supervise the generator, which is called MCMC teaching. Further
research in [25] shows that this cooperative learning method could also
provide a good start point for adversarial models with small computation
overhead. Additionally, the model could also be extended for image-to-image
translation [8] with two pairs of descriptor and generator or used in saliency
prediction [26] by introducing a conditional latent variable model.

#### Progressive Learning

The proposed idea of progressive cooperative learning is closely connected to
the research conducted by [27], which involves the incremental growth of a
single EBM. The multi-grid EBM framework [28], trains a series of EBMs
simultaneously at various resolutions. The sampling process is conducted
sequentially, starting from low-resolution and gradually progressing to higher
resolutions, leveraging the lower resolution as a foundation for subsequent
higher-resolution sampling. In contrast, our method, which combines the growth
of an EBM with three mapping networks, introduces a more challenging and
complex progressive learning strategy. It is important to note that while
there are several progressive learning frameworks based on Generative
Adversarial Networks (GANs), our approach falls within the domain of energy-
based learning. We need to carefully consider MCMC sampling when progressively
expanding the energy function, as it plays a crucial role in both the bottom-
up energy mapping and top-down image generation.

##  3 Proposed Framework

Figure 1: Diagram of cooperative learning for multi-domain image-to-image
translation. The framework consists of a style generator, a style encoder, a
translator and a descriptor. Given a input source image, the translator can
transform it into a target domain, which is specified by a style code. The
style code can be obtained by sampling from the domain-specific style
generator or extracted from a reference image by the style encoder. The
descriptor is a multi-domain image distribution, which plays the role of
guiding the translation such that the translated images can match the observed
images in the target domain in terms of statistical property. All components
are trained simultaneously in a cooperative learning scheme. The descriptor
learns from the multi-domain training images by maximizing the data
likelihood, while utilizing MCMC teaching to guide the training of the
translator, style encoder, and style generator.

Suppose we have unpaired images from multiple domains A, B, C, $\cdots$ with
some shared high-level features, such as expressions in face images, our
target is to learn a conditional generative model that maps an image into a
target domain, which could be same as the source domain, with specific
features. To achieve this, we propose a generative model that consists of four
components, i.e., descriptor, style encoder, style generator, and translator.
The latter three can form a diversified translator, which is trained with the
descriptor in a cooperative learning manner. Let $x$ be an observed image and
$y$ be its domain label. We also use $y^{\prime}$ to denote the label of
target domain.

###  3.1 Multi-Domain Image Descriptor

The multi-domain image descriptor is a multi-head energy-based model that
specifies the probability distribution of each domain by

$p_{y}(x;\theta)\propto\exp[D_{y}(x;\theta)],$ (1)

where $\theta$ are parameters of the multi-head energy function $D$. For
notation simplicity, we use $D_{y}(\cdot)$ to denote the negative energy for
domain $y$. The descriptor are learned by multi-domain maximum likelihood
estimation, which is equivalent to minimizing the Kullback-Leibler (KL)
divergence between the data distribution $p_{\text{data}}(x,y)$ and the model
$p_{y}(x;\theta)$. The gradient of the objective for learning the descriptor
is given by

$\nabla_{\theta}\mathcal{L}_{\text{ebm}}(\theta)=-\mathbb{E}_{p_{\text{data}}(x,y)}\\{\nabla_{\theta}D_{y}(x;\theta)-\mathbb{E}_{p_{y}(x^{\prime};\theta)}\left[\nabla_{\theta}D_{y}\left(x^{\prime};\theta\right)\right]\\},$
(2)

where $\mathbb{E}_{p_{y}(x^{\prime};\theta)}$ denotes the expectation with
respect to the EBM and we use $x^{\prime}$ in order to distinguish the random
variable $x$ in $\mathbb{E}_{p_{\text{data}}(x,y)}$ in the same equation.
Suppose we observe a batch of training examples $\\{(x_{i},y_{i})\\}_{i}^{n}$,
which is assumed to be from $p_{\text{data}}(x,y)$. The gradient in Eq.(2) can
be approximated by

$\nabla_{\theta}\mathcal{L}_{\text{ebm}}(\theta)\approx\nabla_{\theta}\left[\frac{1}{n}\sum_{i=1}^{n}D_{y_{i}}(x_{i};\theta)-\frac{1}{n}\sum_{i=1}^{n}D_{y_{i}}(\tilde{x}_{i};\theta)\right],$
(3)

where for each observed domain $y_{i}$, we use Langevin dynamics to obtain the
corresponding synthesized example $\tilde{x}_{i}$ as a sample from
$p_{y_{i}}(x;\theta)$. With a specified step size $\delta$, Langevin dynamics
is performed by iterating the follow step

$\tilde{x}_{\tau+1}=\tilde{x}_{\tau}+\delta\nabla_{x}D_{y}\left(\tilde{x}_{\tau};\theta\right)+\sqrt{2\delta}U_{\tau},\hskip
5.69054ptU_{\tau}\sim\mathcal{N}(0,I),$ (4)

where $\tau$ indexes time step and $\tilde{x}_{\tau=0}$ is initialized by the
output of a style-controlled image-to-image translator, which is presented in
Section 3.2. A good initialization improves the efficiency of Langevin
dynamics. To stabilize the EBM training, we also add an $l_{2}$ regularization
on the energy outputs of both training examples and synthesized examples,
which is

$\mathcal{L}_{\text{energy}}(\theta)=\frac{1}{n}\sum_{i=1}^{n}\|D_{y_{i}}(x_{i};\theta)\|^{2}+\frac{1}{n}\sum_{i=1}^{n}\|D_{y_{i}}(\tilde{x}_{i};\theta)\|^{2}.$
(5)

Figure 2: An illustration of the progressive strategy for the style encoder
$E$, translator $T$, and descriptor $D$. Boxes in dark grey represent well-
trained modules at resolution level $s-1$, while blocks in light gray
represent the newly added parameters at the current resolution level $s$. The
expansion of the model involves removing some incompatible parameters
(depicted as dark grey boxes with dashed boundaries) and adding new parameters
(depicted as light grey boxes). The output of the module that needs to be
removed and the output of the module that needs to be added are fused using a
transition factor $\omega$, This factor starts from 0 and gradually increases
to 1, controlling the percentage of contribution from the old and new modules.
Left: style encoder. Middle: style-controlled image-to-image translator.
Right: descriptor.

###  3.2 Diversified Image Generator

#### Multi-Domain Style Generator

Given a latent variables $z$ and a domain label $y$, the multi-domain style
generator can produce a domain-specific style code $c$ by

$c_{y}=G_{y}(z;\beta)+\epsilon,\hskip
5.69054pt\epsilon\sim\mathcal{N}(0,I),z\sim\mathcal{N}(0,I),$ (6)

where $\epsilon$ is an observation residual and $z$ follows a Gaussian prior
distribution. $G$ is a multilayer perceptron (MLP) with multiple output
branches to produce style codes for multiple domains. The distribution of
style code $c$ conditioned on a domain $y$ is given by $p_{y}(c;\beta)=\int
p_{y}(c|z;\beta)p(z)dz$, which is more informative than the prior distribution
$p(z)$ to capture the underlying style space. The domain-specific style code
$c_{y}$ is directly used in the translator, which is presented in Section 3.2,
for specifying the style and the target domain of the translated image.

####  Style Encoder

The style encoder $E$ is a multi-head bottom-up network that takes as input an
image $x$ and its corresponding domain label $y$ and then outputs a domain-
specific style code $c=E_{y}(x;\phi)$, where $\phi$ are parameters and
$E_{y}(\cdot)$ denotes the output of $E$ that corresponds to domain $y$.

#### Style-Controlled Image-to-Image Translator

To achieve a one-to-many translation between domains, we build a style-
controlled image-to-image translator. It is a conditioned encoder-decoder $T$
that takes as input a source reference image $x$ and a domain-specific style
code $c_{y}$ and outputs a translated image in target domain $y$, which is
given by

$x_{y}=T(x,c_{y};\alpha)+\epsilon,\hskip
5.69054pt\epsilon\sim\mathcal{N}(0,I),c_{y}\sim p_{y}(c;\beta),$ (7)

where $\alpha$ is the parameters of the neural network $T$. The randomness in
the translated image, when given a reference image and the target domain,
arises from the stochastic nature of the style codes, which follows a
distribution defined by the style generator $p_{y}(c;\beta)$. The translator
$T$ and the style generator $G$ forms a diversified translator. They are
trained by the MCMC teaching loss:

$\mathcal{L}_{\text{teach}}(\alpha,\beta)=\mathbb{E}_{z,y,x}[\|\tilde{x}_{z,y,x}-T(x,G_{y}(z;\beta);\alpha)\|^{2}],$
(8)

where $\tilde{x}_{z,y,x}$ denotes the Langevin synthesis from the descriptor,
which is initialized by the output of $T(x,G_{y}(z;\beta);\alpha)$. That is,
we set $\tilde{x}_{z,x,y,\tau=0}\leftarrow T(x,G_{y}(z;\beta)$ for Langevin
dynamics in Eq.(4) to revolve $\tilde{x}_{z,y,x}$. Let
$M_{\theta}q_{\alpha,\beta}(x)$ be the marginal distribution obtained by
running Markov transition $M_{\theta}$ from $q(x;\alpha,\beta)$. At learning
step $t+1$, the gradient of the MCMC teaching loss in Eq.(8) is the gradient
of
$\text{KL}(M_{\theta^{(t)}}q_{\alpha^{(t)},\beta^{(t)}}||q_{\alpha,\beta})$,
where $q_{\alpha,\beta}$ seeks to be the stationary distribution of
$M_{\theta}$, i.e., minimizing $\text{KL}(p_{\theta}||q_{\alpha,\beta})$. The
effects of the MCMC teaching loss include: (i) $q$ can chase $p$ toward
$p_{\text{data}}$ for MLE; (ii) $q$ can serve as a good MCMC initializer for
$p$ for efficient MCMC sampling. To ensure diverse translator outputs, we
regularize $T$ by minimizing the negative diversity sensitive

$\mathcal{L}_{\text{diverse}}(\alpha)=-\mathbb{E}_{z_{1},z_{2},y,x}[\|T(x,G_{y}(z_{1};\beta);\alpha)-T(x,G_{y}(z_{2};\beta);\alpha)\|_{1}].$
(9)

Since the translator is learned from unpaired data domains, to ensure the
translated image $T(x,c;\alpha)$ to preserve the domain-invariant features of
the source image $x$, we adopt the cycle consistency loss

$\mathcal{L}_{\text{cycle}}(\alpha)=\mathbb{E}_{z,y,x,y^{\prime}}[\|x-T(T(x,G_{y^{\prime}}(z;\beta);\alpha),E_{y}(x;\phi);\alpha)\|_{1}].$
(10)

To ensure any style code that is applied to the translated image can be
retrieved back from the translated image by the style encoder, we also have a
style code reconstruction loss

$\mathcal{L}_{\text{style}}(\alpha,\phi)=\mathbb{E}_{z,y^{\prime},x}[\|G_{y^{\prime}}(z;\beta)-E_{y^{\prime}}(T(x,G_{y^{\prime}}(z;\beta);\alpha);\phi)\|_{1}].$
(11)

To further stabilize the cooperative training and accelerate the MCMC teaching
effect, we propose to add the following energy-based regularization on the
translator,

$\mathcal{L}_{\text{mode}}(\alpha,\beta)=\mathbb{E}_{z,y^{\prime},x}[D_{y^{\prime}}(T(x,G_{y^{\prime}}(z;\beta);\alpha);\theta)],$
(12)

which can shift the translator mapping toward the low energy modes of the
energy function.

###  3.3 Cooperative Learning of Image Descriptor and Image Generator

Our full objective function of the image descriptor is
$\mathcal{L}_{\text{descriptor}}=\mathcal{L}_{\text{ebm}}+\lambda_{\text{energy}}\mathcal{L}_{\text{energy}}$
and the full objective function of the image generator is
$\mathcal{L}_{\text{generator}}=\mathcal{L}_{\text{teach}}+\lambda_{\text{diverse}}\mathcal{L}_{\text{diverse}}+\lambda_{\text{cycle}}\mathcal{L}_{\text{cycle}}+\lambda_{\text{style}}\mathcal{L}_{\text{style}}+\lambda_{\text{mode}}\mathcal{L}_{\text{mode}}$,
where $\lambda_{\text{energy}}$, $\lambda_{\text{diverse}}$,
$\lambda_{\text{cycle}}$, $\lambda_{\text{style}}$, and
$\lambda_{\text{mode}}$ are hyperparameters. At each learning iteration, the
cooperative learning algorithm alternates the following steps: (1) Generate an
initial translated image via $\hat{x}=T(x,G_{y}(z))$; (2) Revise $\hat{x}$ by
Langevin dynamics in Eq.4 to obtain $\tilde{x}$; (3) Update the parameters
$\theta$ of descriptor by minimizing $\mathcal{L}_{\text{descriptor}}$; (4)
Update the parameters $\alpha,\phi,\beta$ of generator by minimizing
$\mathcal{L}_{\text{generator}}$.

Input: Multi-resolution data $\\{(x^{(s)}_{i},y^{(s)}_{i})$,
$i=1,...,N;s=1,...,S\\}$

Output: Model $E^{(S)},T^{(S)},D^{(S)},G$

$E^{(0)}\leftarrow\varnothing,D^{(0)}\leftarrow\varnothing,T^{(0)}\leftarrow\varnothing$

for $s=1,\cdots,S$ do

$m\leftarrow 0$

if $s==1$ then

$\omega\leftarrow 1$

else

$\omega\leftarrow 0$

end if

$E^{(s,\omega)}\leftarrow\text{expand}(E^{(s-1)})$

$D^{(s,\omega)}\leftarrow\text{expand}(D^{(s-1)})$

$T^{(s,\omega)}\leftarrow\text{expand}(T^{(s-1)})$

while ($m\leq N$) do

Sample $(x,y)$ and $y^{\prime}$

Sample $z\sim\mathcal{N}(0,I)$

$c\leftarrow E_{y}^{(s,\omega)}(x)$ or $c\leftarrow G_{y^{\prime}}(z)$

$\hat{x}\leftarrow T^{(s,\omega)}(x,c)$

Revise $\hat{x}$ to obtain $\tilde{x}$ by a $K$-step Langevin dynamics in Eq.
(4).

Update descriptor $D^{(s,\omega)}$ with $\mathcal{L}_{\text{descriptor}}$

Update translator $\\{E^{(s,\omega)},T^{(s,\omega)},G\\}$ with
$\mathcal{L}_{\text{translator}}$

$m\leftarrow m+n^{(s)}$

if $s\neq 1$ then

$\omega\leftarrow\min(1,m/N)$

end if

end while

end for

Algorithm 1 Progressive Cooperative Learning

###  3.4 Progressive Cooperative Learning

The update of both descriptor and generator relies on the cooperative
generation of MCMC synthesized examples, denoted as $\tilde{x}$. To
significantly improve training efficiency, we propose a progressive learning
strategy for our cooperative learning framework. The algorithm gradually
enhances the model resolution from low to high, while maintaining cooperative
learning across all components at each resolution. The underlying motivation
behind this strategy is that learning and sampling from a low resolution data
domain is much more efficient. By leveraging a pre-trained low resolution
model as a foundation, we can efficiently learn the next scale of the model,
rather than starting from scratch. When expanding the current model to the
next scale, each component’s network structure undergoes modifications. New
layers are added to handle higher resolution image inputs or outputs, while
incompatible old layers are removed. The newly added layers are trained
together with the remaining parameters. To ensure a smooth transition and
prevent gradient exposure due to the addition of expanding blocks in each
component, we propose to retain partial effects of the parameters that need to
be removed while incorporating the effects of the newly added parameters.
Throughout each resolution of learning, the impact of the removed parameters
gradually diminishes until it becomes zero. Figure 2 illustrates the expanding
strategy of each component at every level of resolution. Here, $\omega$
represents a transition factor that starts from 0 and increase to 1,
controlling the percentage of effects from the parameters to be removed
(depicted as dark grey boxes with dashed boundaries) and the parameters to be
added (depicted as light grey boxes). For a complete description of the
proposed progressive cooperative learning algorithm, please refer to Algorithm
1.

##  4 Experiment

###  4.1 Experiment Settings

#### Dataset and evaluation metrics

To demonstrate the performance of our proposed multi-domain image-to-image
translation framework, we test it on the CelebA-HQ [9] and AFHQ [2] datasets
and compare it with several baselines. We use M and F to refer the domains of
male and female in CelebA-HQ, and C, D and W to refer to the domains of cat,
dog, and wild animals in AFHQ. We only use the images and the corresponding
domain labels from the datasets in our experiments. We evaluate the quality of
translated images using the Fréchet inception distance (FID) [29], which is a
metric to measure the distance between the population of translated images and
the population of original images in the target domain. A small FID is desired
to indicate that the translated distribution is very close to the target
distribution.

#### Hyper-parameters and network architectures

We use bottom-up convolutional neural networks to parameterize the descriptor
and style encoder and use a 8-layer multilayer perceptron (MLP) for the style
generator. The Translator utilizes an encoder-decoder architecture with AdaIN
[30] for style control in the decoding part. As to the progressive learning,
we start from training our model with a resolution of 64$\times$64, and then
gradually expand the model to resolutions of 128$\times$128 and 256$\times$256
in sequence. We employ 16 Langevin steps for MCMC sampling at the initial
stage of the algorithm, and then reduce the number of steps by 4 after each
stage of expansion. The hyperparameters of $\lambda_{\text{energy}}$,
$\lambda_{\text{diverse}}$, $\lambda_{\text{cycle}}$,
$\lambda_{\text{style}}$, and $\lambda_{\text{mode}}$ are set to be 1. We show
the details of the network structures in Figure 3, in which the modules
depicted with dotted lines indicate that they will be omitted during model
expansion. Since the style generator does not involve changes in image
resolution, its architecture remains unchanged as progressive learning
progresses. For each DownBlock/UpBlock module at each expansion stage $s$, the
input and output channels are set to $\text{max}(512,2^{5+s})$ and
$\text{max}(512,2^{6+s})$, respectively. In the fully-connected (FC) Layer of
the style generator, the input channel of the top FC Layer is set to 16, while
the remaining FC Layers have an input channel of 512. The size of the style
code is consistently set to 64 for all experiments.

Figure 3: Illustration of network architectures for the descriptor,
translator, style Generator, and style Encoder in the proposed progressive
cooperative learning framework for multi-domain image-to-image translation.
The modules depicted with dotted lines indicate that they will be omitted
during model expansion.

###  4.2 Diverse Image Generation

In this experiment, we use style codes that are randomly sampled from the
style generator to generate diverse translated images. Examples of generation
results on the CelebA-HQ dataset can be seen in Figure 4. For each source
image shown in the first row, we generate multiple outputs using random
Gaussian noise. The qualitative results verify the diversity of the translated
results from a source input image. We observe that, given a source image, our
model can not only generate diverse translated images but also produce high-
quality images that obtain the same attribute (e.g., expression) from the
source image. Figure 5 exhibits qualitative results for the AFHQ dataset. It
is evident that our model successfully produces diverse, reasonable, and
realistic translated images.

Figure 4: Qualitative results of diverse image generation on CelebA-HQ
dataset. Each column displays one example of one-to-many image generation. The
first row displays source images. The rest four rows show different translated
images, which are obtained by using four style codes randomly generated by the
style generator. The style generator produces style codes by randomly sampling
from Gaussian distribution. Figure 5: Qualitative results of diverse image
generation on AFHQ dataset. Each column show one example of one-to-many image
generation. The first row displays source images and the rest six rows show
different translated images generated by using different style codes sampled
from the style generator.

###  4.3 Image-to-Image Translation using reference image

In our image-to-image translation approach, we can incorporate a reference
image. To begin with, we adopt the style encoder to extract the style code
from the provided reference image. Subsequently, we apply the style code to
the translator network, enabling the generation of translated images based on
the desired style conveyed by the reference image. Figure 6 show some
qualitative results obtained on dataset CelebA-HQ, where we take images in the
first row as source images and images in the first column as reference images.
The translation results are shown in the middle. Comparing results displayed
in each row, we can observe that the human face in the source image can be
clearly changed into the same gender and appearance of the face in the
reference image, while keeping the facial expression consistent with that in
the source domain. Figure 7 presents examples demonstrating the usage of
reference images for image-to-image translation in the AFHQ dataset. These
results align with our observations and findings. The translated images
exhibit similar appearance styles as the reference image, indicating that they
successfully capture the desired target domain properties. Simultaneously, the
translated images preserve the geometry aspects, including shapes and poses,
from the source images. This indicates that our approach effectively maintains
the essential characteristics of the source images while incorporating the
desired style from the reference image.

To verify the cycle consistency of our model, we conduct experiments for cycle
image-to-image translation. Firstly, we perform an image-to-image translation
from the source domain to the target domain, guided by a reference image.
Then, we reverse the translation by taking the translated image as the new
source image and translating it back to the original domain. We assess the
quality of the cycle consistency through qualitative results, displayed in
Figure 8 for the CelebA-HQ dataset and Figure 9 for the AFHQ dataset. The
qualitative results clearly demonstrate that our model exhibits excellent
cycle consistency properties. It effectively preserves the geometric features
of the source image while applying the appearance properties from the
reference image during translation in both directions. By comparing the images
in the first and fourth rows, we can see that our algorithm can maintain a
high level of synthesis quality and provide reasonable results in a sequential
usage of the model. Our model can ensure a high level of fidelity and
consistency in the translation process.

We also illustrate the progressive learning process by showcasing translated
images at various resolutions. Figure 10 presents the results for the CelebA-
HQ dataset, while Figure 11 showcases the results for the AFHQ dataset. Each
column in the figures corresponds to translated images generated by models at
different stages of resolution, using a specific source image and reference
image. As observed, the generated images become more detailed as the model
resolution increases.

Figure 6: We show the translated images with style codes generated from the Style Encoder and reference images in CelebA-HQ dataset in this picture. The source images and reference images are put in the first row and first column. We could see that the face has successfully translated into target domains with consistency in expression. Figure 7: Results for reference generation on AFHQ dataset. The source images (first row) take in the style codes generated by the Style Encoder and reference images (first column) and translates into target domains. We could see the animals’ faces have successfully translated into target domains with consistency in expression. Figure 8: Qualitative results of cycle consistency check on CelebA-HQ dataset. Each column corresponds to a specific example. The first row displays an image from the source domain A, while the second row shows a reference image from domain B. The third row exhibits the translated image obtained by translating the source image to the target domain specified by the reference image. The fourth row presents the reversed translation results obtained by using the image from the third row and treating the image from the first row as the reference image. Figure 9: Qualitative results of cycle consistency check on AFHQ dataset. Each column shows one example. The first row displays an image from the source domain A. The second row showcases a reference image from domain B, providing the desired style or specifying the target domain. The third row exhibits the translated image, which is obtained by transforming the source image to the target domain specified by the reference image. The fourth row presents the results of the reversed translation, achieved by utilizing the image from the third row and considering the image from the first row as the reference image. Table 1: Evaluation on CelebA-HQ dataset. Dim |  Method | Reference |  | Diverse
---|---|---|---|---
|  |  M$\rightarrow$F |  F$\rightarrow$M |  Avg |  |  M$\rightarrow$F |  F$\rightarrow$M |  Avg
256 |  MUNIT[31] |  - |  - |  107.1 |  |  - |  - |  31.4
|  DRIT[32] |  - |  - |  53.3 |  |  - |  - |  52.1
|  MSGAN[33] |  - |  - |  39.6 |  |  - |  - |  61.4
|  StarGAN2[2] |  19.8 |  28.0 |  23.8 |  |  10.0 |  17.6 |  13.8
|  ours |  19.9 |  30.9 |  25.4 |  |  22.1 |  32.3 |  27.2
128 |  CycleCoop[8] |  - |  - |  - |  |  122.4 |  131.9 |  127.1
|  StarGAN2[2]* |  21.1 |  27.2 |  24.2 |  |  11.9 |  17.6 |  14.8
|  ours |  14.5 |  22.9 |  18.7 |  |  13.7 |  23.0 |  18.4
64 |  CycleCoop[8] |  - |  - |  - |  |  129.8 |  124.5 |  127.2
|  StarGAN2[2] |  18.8 |  22.2 |  20.5 |  |  10.2 |  13.3 |  11.8
|  ours |  10.8 |  19.3 |  15.1 |  |  10.5 |  18.2 |  14.3
Table 2: Evaluation on AFHQ dataset for diverse image generation. Dim |  Method |  D$\rightarrow$C |  W$\rightarrow$C |  C$\rightarrow$D |  W$\rightarrow$D |  C$\rightarrow$W |  D$\rightarrow$W |  Avg
---|---|---|---|---|---|---|---|---
256 |  MUNIT[31] |  - |  - |  - |  - |  - |  - |  41.5
|  DRIT[32] |  - |  - |  - |  - |  - |  - |  95.6
|  MSGAN[33] |  - |  - |  - |  - |  - |  - |  61.4
|  StarGAN2[2] |  6.4 |  7.1 |  36.2 |  30.4 |  9.1 |  8.6 |  16.3
|  ours |  15.8 |  15.4 |  42.5 |  42.4 |  28.4 |  28.5 |  28.8
128 |  StarGAN2* |  7.1 |  6.6 |  32.9 |  29.6 |  10.2 |  9.4 |  16.0
|  ours |  12.8 |  13.3 |  36.5 |  35.8 |  11.5 |  11.6 |  20.2
64 |  StarGAN2* |  8.6 |  7.0 |  34.6 |  28.9 |  15.3 |  17.8 |  18.7
|  ours |  13.6 |  13.8 |  33.2 |  32.5 |  12.7 |  12.6 |  19.7
Table 3: Evaluation on AFHQ dataset for image translation using reference images. Dim |  Method |  D$\rightarrow$C |  W$\rightarrow$C |  C$\rightarrow$D |  W$\rightarrow$D |  C$\rightarrow$W |  D$\rightarrow$W |  Avg
---|---|---|---|---|---|---|---|---
256 |  MUNIT[31] |  - |  - |  - |  - |  - |  - |  223.9
|  DRIT[32] |  - |  - |  - |  - |  - |  - |  114.8
|  MSGAN[33] |  - |  - |  - |  - |  - |  - |  69.8
|  StarGAN2[2] |  6.3 |  7.3 |  39.9 |  33.1 |  16.4 |  15.6 |  19.8
|  ours |  11.2 |  10.7 |  33.1 |  32.8 |  8.5 |  8.2 |  17.4
128 |  CycleCoop |  93.5 |  210.4 |  102.8 |  183.8 |  182.2 |  204.2 |  162.8
|  StarGAN2[2] |  7.3 |  9.0 |  37.2 |  33.6 |  8.9 |  10.4 |  17.8
|  ours |  9.6 |  9.6 |  27.0 |  25.8 |  7.4 |  6.7 |  14.3
64 |  CycleCoop |  267.6 |  259.0 |  194.5 |  226.5 |  223.2 |  212.9 |  230.6
|  StarGAN2[2] |  7.3 |  9.1 |  35.3 |  31.0 |  18.6 |  23.6 |  20.8
|  ours |  11.7 |  10.5 |  24.6 |  24.0 |  7.1 |  7.5 |  14.2

###  4.4 Quantitative Comparison

We also compare the results of our translation results quantitatively by using
style codes from Style Encoder by randomly selecting reference image in
different domains or from Style Generator through sampling from Gaussian
distribution. For each source image in the validation dataset, we obtain ten
translated images for each target domain to compute the FID. Results are shown
in Table 1 for CelebA-HQ and Table 2& 3 for AFHQ. We compare our model with
other baselines, including GAN-based method, StarGAN v2 [2], and energy-based
method, CycleCoop [8]. We could see that our model could significantly out
perform existing cooperative learning methods with additional ability of
guidance by reference images and reach comparable performance with GAN-based
methods.

Figure 10: Progressive learning on CelebA-HQ dataset. Translated images are
generated by models learned at resolutions of $64\times 64$, $128\times 128$,
$256\times 256$, respectively. The same source image and reference image are
utilized for generating images at each column. Figure 11: Progressive learning
on AFHQ dataset. Translated images are generated by models learned at
resolutions of $64\times 64$, $128\times 128$, $256\times 256$, respectively.
The same source image and reference image are utilized for generating images
at each column.

###  4.5 Ablation Study

We conduct an ablation study to evaluate the importance of each individual
component proposed in our paper. In Table 4, we report the model performance
in terms of FID by removing different key loss term (including
$\mathcal{L}_{\text{diverse}},\mathcal{L}_{\text{cycle}},\mathcal{L}_{\text{energy}},\mathcal{L}_{\text{mode}}$)
from our objective function in our framework. We train our model in a
64$\times$64 resolution setting on datasets CelebA-HQ and AFHQ without using
the progressive learning strategy. We show results of image translation using
style codes obtained from both style encoder and style generator and report
average performance in Table 4. NA means that the model fails in learning and
can not generate meaningful results. We can see that the newly added
regularization strategies for the descriptor and the translator, i.e.,
$\mathcal{L}_{energy}$ and $\mathcal{L}_{mode}$, are essential for stabilizing
the cooperative training. Especially, the energy-based regularization loss
$\mathcal{L}_{mode}$ plays an important role to ensure that the translator can
quickly catch up with the descriptor toward the data distribution during the
cooperative training because we learn the translator so that its output can
minimize the energy function. The $\mathcal{L}_{energy}$ is useful to obtain
performance gain by limiting the magnitude of the energy values. Also, we can
find that the performance drops significantly when removing the cycle-
consistency loss $\mathcal{L}_{cycle}$, which has proved to be a key objective
for unpaired cross-domain image translation task.

Table 4: Ablation Study on CelebA-HQ and AFHQ datasets in 64$\times$64 resolution. Removed Item | CelebA |  | AFHQ |  Avg
---|---|---|---|---
|  Reference |  Latent |  |  Reference |  Latent
baseline |  15.1 |  14.3 |  |  14.2 |  19.7 |  17.0
Remove $\mathcal{L}_{diverse}$ |  16.3 |  17.1 |  |  36.4 |  35.2 |  26.3
Remove $\mathcal{L}_{cycle}$ |  111.0 |  127.3 |  |  NA |  NA |  119.2
Remove $\mathcal{L}_{energy}$ |  134.5 |  40.7 |  |  208.5 |  97.6 |  120.3
Remove $\mathcal{L}_{mode}$ |  NA |  NA |  |  277.2 |  217.6 |  247.4

##  5 Conclusion

We present MD-CoopNets, a novel approach that combines energy-based learning,
MCMC sampling, cooperative learning, and progressive learning for unpaired
multi-domain image-to-image translation. Our method includes a multi-head
energy-based model as a descriptor, capturing the multi-domain image
distribution, and a diversified image-to-image translator for cross-domain
one-to-many mapping. To train both the descriptor and translator, we introduce
a multi-domain MCMC teaching algorithm. Additionally, we propose progressive
learning to enhance the scalability and efficiency of our model. Experimental
results demonstrate that our approach achieves comparable performance to
adversarial learning frameworks and sets a new benchmark in energy-based
image-to-image translation methods.

## References

  * [1] Yunjey Choi, Min-Je Choi, Munyoung Kim, Jung-Woo Ha, Sunghun Kim, and Jaegul Choo.  Stargan: Unified generative adversarial networks for multi-domain image-to-image translation.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 8789–8797, 2018.
  * [2] Yunjey Choi, Youngjung Uh, Jaejun Yoo, and Jung-Woo Ha.  Stargan v2: Diverse image synthesis for multiple domains.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 8188–8197, 2020.
  * [3] Jianwen Xie, Yang Lu, Ruiqi Gao, Song-Chun Zhu, and Ying Nian Wu.  Cooperative training of descriptor and generator networks.  IEEE transactions on pattern analysis and machine intelligence, 42(1):27–45, 2018.
  * [4] Jianwen Xie, Zilong Zheng, and Ping Li.  Learning energy-based model with variational auto-encoder as amortized sampler.  In Thirty-Fifth AAAI Conference on Artificial Intelligence, (AAAI)), pages 10441–10451, 2021.
  * [5] Jianwen Xie, Yaxuan Zhu, Jun Li, and Ping Li.  A tale of two flows: Cooperative learning of langevin flow and normalizing flow toward energy-based model.  In International Conference on Learning Representations (ICLR), 2022\.
  * [6] Jing Zhang, Jianwen Xie, Zilong Zheng, and Nick Barnes.  Energy-based generative cooperative saliency prediction.  In Thirty-Sixth AAAI Conference on Artificial Intelligence, (AAAI)), pages 3280–3290, 2022.
  * [7] Jianwen Xie, Zilong Zheng, Xiaolin Fang, Song-Chun Zhu, and Ying Nian Wu.  Cooperative training of fast thinking initializer and slow thinking solver for conditional learning.  IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI), 44(8):3957–3973, 2022.
  * [8] Jianwen Xie, Zilong Zheng, Xiaolin Fang, Song-Chun Zhu, and Ying Nian Wu.  Learning cycle-consistent cooperative networks via alternating mcmc teaching for unsupervised cross-domain translation.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10430–10440, 2021.
  * [9] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen.  Progressive growing of gans for improved quality, stability, and variation.  arXiv preprint arXiv:1710.10196, 2017.
  * [10] Song Chun Zhu, Yingnian Wu, and David Mumford.  Filters, random fields and maximum entropy (frame): Towards a unified theory for texture modeling.  International Journal of Computer Vision, 27:107–126, 1998.
  * [11] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang.  A tutorial on energy-based learning.  Predicting structured data, 1(0), 2006.
  * [12] Geoffrey E. Hinton.  A practical guide to training restricted boltzmann machines.  In Neural Networks: Tricks of the Trade - Second Edition, pages 599–619. Springer, 2012.
  * [13] Jianwen Xie, Yang Lu, Song-Chun Zhu, and Ying Nian Wu.  A theory of generative convnet.  In Maria-Florina Balcan and Kilian Q. Weinberger, editors, International Conference on Machine Learning (ICML), 2016.
  * [14] Erik Nijkamp, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu.  Learning non-convergent non-persistent short-run mcmc toward energy-based model.  Advances in Neural Information Processing Systems, 32, 2019.
  * [15] Yilun Du and Igor Mordatch.  Implicit generation and generalization in energy-based models.  In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems, (NeurIPS), 2019.
  * [16] Geoffrey E Hinton.  Training products of experts by minimizing contrastive divergence.  Neural computation, 14(8):1771–1800, 2002.
  * [17] Yilun Du, Shuang Li, Joshua B. Tenenbaum, and Igor Mordatch.  Improved contrastive divergence training of energy-based models.  In Proceedings of the 38th International Conference on Machine Learning, ICML, 2021.
  * [18] Ruiqi Gao, Yang Song, Ben Poole, Ying Nian Wu, and Diederik P. Kingma.  Learning energy-based models by diffusion recovery likelihood.  In The ninth International Conference on Learning Representations, ICLR, 2021.
  * [19] Ruiqi Gao, Erik Nijkamp, Diederik P Kingma, Zhen Xu, Andrew M Dai, and Ying Nian Wu.  Flow contrastive estimation of energy-based models.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR, 2020.
  * [20] Taesup Kim and Yoshua Bengio.  Deep directed generative models with energy-based probability estimation.  arXiv preprint arXiv:1606.03439, 2016.
  * [21] Tian Han, Erik Nijkamp, Xiaolin Fang, Mitch Hill, Song-Chun Zhu, and Ying Nian Wu.  Divergence triangle for joint training of generator model, energy-based model, and inferential model.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 8670–8679, 2019.
  * [22] Rithesh Kumar, Sherjil Ozair, Anirudh Goyal, Aaron Courville, and Yoshua Bengio.  Maximum entropy generators for energy-based models.  arXiv preprint arXiv:1901.08508, 2019.
  * [23] Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat.  Vaebm: A symbiosis between variational autoencoders and energy-based models.  In The ninth International Conference on Learning Representations, ICLR, 2021.
  * [24] Will Sussman Grathwohl, Jacob Jin Kelly, Milad Hashemi, Mohammad Norouzi, Kevin Swersky, and David Duvenaud.  No MCMC for me: Amortized sampling for fast and stable training of energy-based models.  In The ninth International Conference on Learning Representations, ICLR, 2021.
  * [25] Yang Zhao, Jianwen Xie, and Ping Li.  Coopinit: Initializing generative adversarial networks via cooperative learning.  arXiv preprint arXiv:2303.11649, 2023.
  * [26] Jing Zhang, Jianwen Xie, Zilong Zheng, and Nick Barnes.  Energy-based generative cooperative saliency prediction.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 3280–3290, 2022.
  * [27] Yang Zhao, Jianwen Xie, and Ping Li.  Learning energy-based generative models via coarse-to-fine expanding and sampling.  In The ninth International Conference on Learning Representations, ICLR, 2021.
  * [28] Ruiqi Gao, Yang Lu, Junpei Zhou, Song-Chun Zhu, and Ying Nian Wu.  Learning generative convnets via multi-grid modeling and sampling.  In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 9155–9164, 2018.
  * [29] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  Advances in neural information processing systems, 30, 2017.
  * [30] Xun Huang and Serge Belongie.  Arbitrary style transfer in real-time with adaptive instance normalization.  In Proceedings of the IEEE international conference on computer vision, pages 1501–1510, 2017.
  * [31] Xun Huang, Ming-Yu Liu, Serge Belongie, and Jan Kautz.  Multimodal unsupervised image-to-image translation.  In Proceedings of the European conference on computer vision (ECCV), pages 172–189, 2018.
  * [32] Hsin-Ying Lee, Hung-Yu Tseng, Jia-Bin Huang, Maneesh Singh, and Ming-Hsuan Yang.  Diverse image-to-image translation via disentangled representations.  In Proceedings of the European conference on computer vision (ECCV), pages 35–51, 2018.
  * [33] Qi Mao, Hsin-Ying Lee, Hung-Yu Tseng, Siwei Ma, and Ming-Hsuan Yang.  Mode seeking generative adversarial networks for diverse image synthesis.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1429–1437, 2019.