# Closed-Loop Unsupervised Representation Disentanglement with
$\beta$-VAE Distillation and Diffusion Probabilistic Feedback

Xin Jin1, *, Bohan Li1, 2, BAAO Xie1,
Wenyao Zhang1, 2, Jinming Liu1, 2, Ziqiang Li1, 2, Tao Yang3, Wenjun Zeng1,
1Ningbo Institute of Digital Twin, Eastern Institute of Technology, Ningbo,
China
2Shanghai Jiao Tong University, Shanghai, China, 3Xi’an Jiaotong University,
Xi’an, China
Corresponding<EMAIL_ADDRESS>
###### Abstract

Representation disentanglement may help AI fundamentally understand the real
world and thus benefit both discrimination and generation tasks. It currently
has at least three unresolved core issues: (i) heavy reliance on label
annotation and synthetic data — causing poor generalization on natural
scenarios; (ii) heuristic/hand-craft disentangling constraints make it hard to
adaptively achieve an optimal training trade-off; (iii) lacking reasonable
evaluation metric, especially for the real label-free data. To address these
challenges, we propose a Closed-Loop unsupervised representation
Disentanglement approach dubbed CL-Dis. Specifically, we use diffusion-based
autoencoder (Diff-AE) as a backbone while resorting to $\beta$-VAE as a co-
pilot to extract semantically disentangled representations. The strong
generation ability of diffusion model and the good disentanglement ability of
VAE model are complementary. To strengthen disentangling, VAE-latent
distillation and diffusion-wise feedback are interconnected in a closed-loop
system for a further mutual promotion. Then, a self-supervised Navigation
strategy is introduced to identify interpretable semantic directions in the
disentangled latent space. Finally, a new metric based on content tracking is
designed to evaluate the disentanglement effect. Experiments demonstrate the
superiority of CL-Dis on applications like real image manipulation and visual
analysis.

##  1 Introduction

Disentangled representation learning (DRL) [31] learns the underlying
explainable factors behind the observed data, where the different latent
factors correspond to different properties, respectively. This matches the way
humans perceive this world that we understand an object from its various
properties (e.g., shape, size, color, etc.) [20, 21, 91, 60, 2]. Thus, DRL is
thought to be one of the possible ways for AI to understand the world
fundamentally, thus achieving Artificial General Intelligence (AGI) [3, 49,
82].

Figure 1: The illustration of disentangled representation learning (DRL)
frameworks: (a) VAE-based, (b) GAN-based, (c) Diffusion-based, and (d) our
proposed Closed-loop approach. We can intuitively see that previous works all
relied on heuristic hand-craft loss constraints to infer disentangled latent,
only our method leverages two kinds of generative branches to build a cycle
system for mutually-promoting automatically disentangled learning.

Most existing DRL methods learn the disentangled representation based on
generative models, such as VAE [45, 30, 4, 47, 42, 7, 18, 43, 69], GAN [25, 8,
77, 54, 36, 98], and even recently popular diffusion [48, 85, 94, 86, 89].
Among them, the VAE-based DRL methods always encourage the disentanglement of
the latent variables by forcing the variational posterior to be closer to the
factorized Gaussian prior (see Figure 1(a)), but they generally need to
balance a hard inherent trade-off between the disentangling ability and
generating quality [30, 7, 42]. As shown in Figure 1(b), the GAN-based DRL
methods tend to use the mutual information as disentanglement constraints, but
they typically suffer from the problem of reconstruction quality due to the
gan-inversion difficulty [81, 94]. Besides, VAE- and GAN-based generative
models are prone to learn the factorized or conditioned statistics bias [32],
which essentially conflicts with the ultimate purpose of a generalizable
disentanglement that makes AI understand the world like humans.

Moreover, given the recently popular diffusion probabilistic models (DPMs)
show superiority in image generation quality [33, 73] and unsupervised
controlling (i.e., learning meaningful representations) of images [64],
researchers have started to investigate the possibility of DRL in DPMs. Asyrp
[48] proposes an asymmetric reverse process to discover the semantic latent
space in frozen pre-trained diffusion models, which reveals that DPMs already
have a semantic latent space and thus more suitable for learning disentangled
semantics. The co-current work of DisDiff [94] shown in Figure 1(c) uses the
gradient field of frozen DPMs to achieve a vector-wise disentangled
representation learning. However, this work applies an external heuristic
hand-craft variant-invariant loss [57, 94] to enforce the disentanglement
processes, resulting in reduced model flexibility and posing challenges for
optimization.

In this paper, we propose to leverage two kinds of generative branches of the
diffusion model and VAE model, to build a cycle system for mutually-promoting
DRL. As shown in Figure 1(d)), we construct an unsupervised Closed-Loop
Disentanglement framework dubbed CL-Dis, which is built upon a diffusion-based
autoencoder (Diff-AE) backbone while combining with a co-pilot of $\beta$-VAE,
to together learn the meaningful disentangled representations. To collaborate
two branches well and facilitate feature disentangling, we encourage Diff-AE
and $\beta$-VAE to complement each other, achieving a win-win effect: we use
the pre-trained latent of VAE to guide a semantic-aware reverse diffusion
process with a distillation loss. Then, the gradually increasing information
capability during the diffusion process is taken as feedback, in turn, to
progressively enforce the inner latent of VAE more disentangling. Next, a
self-supervised Navigation strategy is introduced to clearly identify each
factor’s semantic meanings, where shifting along a disentangled semantic
direction will result in continuously changed generations. Finally, we further
design a label-free metric based on “changes tracking” to quantitatively
measure the disentanglement of learned features, wherein we use optical flow
to reflect image variation degree objectively. Our main contributions are
summarized:

  * •

We build a new unsupervised representation disentanglement framework CL-Dis,
which has a mutually-promoting closed-loop architecture driven by diffusion
and VAE models. They complement each other with a distillation loss and a new
feedback loss, and thus empower a more controllable and stronger DRL
capability.

  * •

We introduce a self-supervised method in CL-Dis to clearly identify each
factor’s semantic meanings by navigating directions in the learned
disentangled latent space, leading to coherent generated variations. This
makes the generation more explainable and fine-controllable.

  * •

We design a new label-free metric to quantitatively measure the
interpretability and disentanglement of learned features, and provide an
example of immediate practical benefit from our work. Namely, we
experimentally show how to exploit our CL-Dis for generative manipulation of
real images and visual discrimination tasks.

##  2 Related Works

###  2.1 Disentangled Representation Learning (DRL)

The concept of DRL was introduced by Bengio et al. [3] in 2013 and was
believed helpful in different tasks in practical applications like image
generation [45, 7, 47, 42, 43, 87, 8, 54], image editing/translation [24, 50,
56, 83], NLP [84, 28, 10] and multimodal applications [34, 78, 97, 90, 59].
Different from image generation or editing, the core of this task is to fully
understand the latent factors of the model to enhance its fine-grained
controllability [7, 47, 42, 43, 87].

Most VAE-based DRL approaches like $\beta$-VAE [30], DIP-VAE [47], FactorVAE
[42], etc., achieve unsupervised disentanglement by the direct constraints on
probabilistic distributions [7, 42, 30]. While straightforward, Locatello et
al. [57] point out that is not enough and emphasize the need for extra
inductive bias. Similarly, Burgess [4] proposes to progressively increase the
information capacity of the latent code in $\beta$-VAE. Yang et al. [92] use
symmetry properties modeled by group theory as inductive bias. Except for VAE,
there are also some DRL works based on GAN, e.g., leveraging mutual
information [8], self-supervised contrastive regularizer [54], spatial
constriction [98], and combining with other pre-trained generative models [52,
66].

Fast-developing generative diffusion probabilistic models (DPMs) have
accelerated the DRL exploration. LCG-DM [44] uses a vanilla VAE in DPMs to
extract semantics codes for controllable generation, but ignores the
interaction between these two structures. Asyrp [48] discovers the semantic
latent space in frozen pre-trained DPMs for image manipulation. Co-current
DisDiff [94] proposes a vector-wise method to express more information
compared with other scalar-based methods [4, 54, 48]. However, they uses
gradient fields as an inductive bias at all time steps to achieve DRL, which
relies on an external heuristic hand-craft disentangling loss, lacking a self-
driven learning ability and is hard to optimize. In contrast, our CL-Dis
builds a closed-loop DRL framework with a mutual-promoting mechanism driven by
a two-branch complementary diffusion and $\beta$-VAE architecture, encouraging
stronger disentangled feature auto-learning.

###  2.2 Knowledge Distillation and Feature Feedback

Our closed-loop CL-Dis is bridged by these two techniques. Knowledge
distillation (KD) is originally developed for model compression and
acceleration [11, 26, 80], or used in transfer learning [38, 1]. Feature
feedback [63, 15, 41, 14, 58] enhances the robustness of learning systems by
using explainable information, which is gaining more and more attention in AI
alignment [37], RLHF [46, 13], and embodied AI [67, 17]. For instance,
ControlVAE [69] achieves much better reconstruction quality with a variant of
the proportional-integral-derivative (PID) control, using the output KL-
divergence as feedback. Following this idea, we explore taking the strong
generation capability of diffusion models as information-increasing feedback
to automatically and adaptively enhance disentangled feature learning.

###  2.3 Latent Semantics Discovery (LSD)

This task emerges to move around the latent code (of GAN [79, 12, 75],
Diffusion [48], etc), such that only one factor varies during the traversal
generation. They mostly rely on human annotations (i.e., segmentation masks,
attribute categories, 3D priors, and text descriptions) to define the semantic
labels [22, 71, 62, 53, 9, 72]. Unsupervised LSD [79, 70, 74] often learn a
set of directions or a classifier/matrix (e.g., Hessian, Jacobian
regularization, etc.) to identify latent semantics. Wu et al. [85] find
semantics in the stable diffusion model by partially changing the text
embeddings. However, these methods are not strictly equivalent to DRL, most of
them just rely on already disentangled (text) features to get coherent
traversal generation results.

##  3 Background and Motivation

###  3.1 Understanding disentangling in $\beta$-VAE

Due to the poor disentanglement of vanilla VAE on complex data, $\beta$-VAE
[30] adds explicit inductive bias to strengthen the independence constraint of
the variational posterior distribution $q_{\phi}(\mathbf{z}|\mathbf{x})$ with
a ${\beta}$ penalty coefficient in ELBO:

$\displaystyle\mathcal{L}(\theta,\phi)=$
$\displaystyle\mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})}\big{[}\log
p_{\theta}(\mathbf{x}|\mathbf{z})\big{]}-\beta
D_{KL}\big{(}q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z})\big{)},$
(1)

when ${\beta}$=1, ${\beta}$-VAE degenerates to the original VAE, and larger
${\beta}$ encourages more disentangled representations but harms the
performance of reconstruction. Unfortunately, it is practically intractable to
obtain the optimal $\beta$ to balance the trade-off between reconstruction
quality and the disentangling capability [47, 7]. Thus, Burgess et al. [4]
understand ${\beta}$-VAE from the perspective of information bottleneck
theory, and propose improving disentangling in ${\beta}$-VAE with controlled
information capacity increase:

$\displaystyle\mathcal{L}(\theta,\phi)=$
$\displaystyle\mathbb{E}_{q_{\phi}(\mathbf{z}\mid\mathbf{x})}\log
p_{\theta}(\mathbf{x}|\mathbf{z})-\beta\big{|}D_{KL}\big{(}q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z})\big{)}-C\big{|},$
(2)

where $C$ gradually increases from $0$ to a value large enough during the
training to guarantee the expressiveness of latent representations, thus
improving disentangling. However, such a hand-craft increasing strategy is
hard to control and might hurt disentangling performance, especially for
various scenarios on different datasets. Thus, in our work, we learn from this
idea but upgrade it by exploiting the essence of the Diffusion Probabilistic
Model (DPM) in automatically increasing information capacity [61, 51] for
disentangling enhancement.

###  3.2 DPM and Diffusion Autoencoder (Diff-AE)

Here, we first take DDPM [33] as a vanilla DPM example for illustration, which
adopts a sequence of fixed variance distributions $q(x_{t}|x_{t-1})$ as the
noise-adding forward process to collapse the image distribution $p(x_{0})$ to
$\mathcal{N}(0,I)$:

$q(x_{t}|x_{t-1})=\mathcal{N}(x_{t};\sqrt{1-\beta_{t}}x_{t-1},\beta_{t}I).$
(3)

The reverse process is fitting by other distributions of $\theta$, and it is a
progressively increasing process of information capacity (i.e., entropy
decreasing) as we mentioned before:

$p_{\theta}(x_{t-1}|x_{t})=\mathcal{N}(x_{t};\mu_{\theta}(x_{t},t),\sigma_{t}I),$
(4)

where $\mu_{\theta}(x_{t},t)$ is parameterize by a Unet
$\epsilon_{\theta}(x_{t},t)$. Its optimization is to minimize the variational
upper bound of negative log-likelihood
$\mathcal{L}_{\theta}=\mathop{\mathbb{E}}_{x_{0},t,\epsilon}\|\epsilon-\epsilon_{\theta}(x_{t},t)\|$.

Furthermore, Diff-AE [64] is proposed for representation learning based on DPM
and Autoencoder. It uses a learnable semantic encoder for discovering the
high-level meaningful semantics $\mathbf{z_{sem}}=E_{sem}(x_{0})$, and a DPM-
based decoder $p(x_{t-1}|x_{t},\mathbf{z_{sem}})$ that is conditioned on
$\mathbf{z_{sem}}$ for reconstruction. However, the latent representation
$\mathbf{z_{sem}}$ learned by Diff-AE does not explicitly respond to the
underlying factors of the data, or say, its disentangling potential has not
been fully unleashed.

Figure 2: Pipeline of the proposed CL-Dis. We explore using (a) diffusion-
based autoencoder (Diff-AE) as a backbone while resorting to (b) $\beta$-VAE
as a co-pilot via knowledge distillation to help $E_{sem}$ extract a semantic
disentangled representation $\mathbf{z_{sem}}$ of an input image $x_{0}$
(which is taken as condition for diffusion-based generation). To further
facilitate feature disentangling, we further take (c) the reverse diffusion
process as an increasing information capacity indicator feedback to, in turn,
enhance the disentanglement of feature $\mathbf{z_{disen}}$ in $\beta$-VAE. As
a result, Diff-AE and $\beta$-VAE are interconnected as a closed-loop cycle
system with VAE-latent distillation and diffusion-wise feedback controlling,
achieving a win-win effect and well-disentangled fine-controllable
representation learning.

##  4 Methodology: Closed-Loop DRL (CL-Dis)

Based on the analysis above, we propose to leverage VAE and diffusion to
unleash their strengths for a mutually promoting DRL. As shown in Figure 2,
our CL-Dis is a dual-branch cycle framework, involving multiple iterative
phases with (a) diffusion autoencoder (Diff-AE), (b) $\beta$-VAE, and the core
(c) knowledge distillation & feature feedback process. In Section 4.1, we
first present the formulation of closed-loop DRL and the overview of CL-Dis.
After that, we introduce different phases of CL-Dis, including Diff-AE pre-
training, $\beta$-VAE pre-training, and knowledge distillation & feature
feedback in Section 4.2. Then, we propose a self-supervised Navigation
strategy to identify disentangled semantic meanings and a new well-designed
DRL metric (Section 4.3) for the final disentanglement measurement.

###  4.1 Problem Formulation and Overview of CL-Dis

Given a practical dataset $\mathcal{D}$ derived by $N$ underlying factors
$\mathcal{C}=\\{1,2,\dots,N\\}$, the target of disentanglement for our CL-Dis
is to learn a semantic encoder $E_{sem}$, for each factor $c\in\mathcal{C}$,
when varying its latent value, its reconstructed generation results will have
the corresponding changes.

In Figure 2, our CL-Dis takes Diff-AE as the backbone, which has exactly a
semantic encoder $E_{sem}$ to encode the raw data $x_{0}$ into the semantic
representations $\mathbf{z_{sem}}$ for conditioning the diffusion reverse
distributions $p(x_{t-1}|x_{t},\mathbf{z_{sem}})$. To empower $E_{sem}$
disentangling capability, we additionally construct a $\beta$-VAE as co-pilot
to help $\mathbf{z_{sem}}$ become a semantic _disentangled_ representation.
Given a dataset $\mathcal{X}$ from a distribution $p_{data}(\mathbf{x})$, a
standard $\beta$-VAE framework [30] is built and $\phi$, $\theta$ parametrize
the VAE encoder $q_{\phi}(\mathbf{z}|\mathbf{x})$ and the decoder
$p_{\theta}(\mathbf{x}|\mathbf{z})$, respectively. The prior $p(\mathbf{z})$
is typically set to the isotropic unit Gaussian $\mathcal{N}(0,1)$, using a
differentiable “reparametrization trick ($z=\mu+\sigma$)” [45]. With the
objective of Eq. 1, we get a preliminary disentangled representations
$\mathbf{z_{disen}}$. Then, we use a distillation loss $\mathcal{L}_{dt}$ to
transfer such disentangling capability to the semantic latent
$\mathbf{z_{sem}}$ of Diff-AE,

$\mathcal{L}_{dt}=D_{KL}\big{(}\hskip 2.84526pt\mathbf{z_{sem}}\hskip
2.84526pt\|\hskip 2.84526pt\mathbf{z_{disen}}\hskip
2.84526pt\big{)}=\sum\mathbf{z_{sem}}\cdot\log(\frac{\mathbf{z_{sem}}}{\mathbf{z_{disen}}}),$
(5)

As discussed before, gradually increasing the information capacity of the VAE
latent with a controlled way could improve disentangling [4]. The reverse
diffusion process of Diff-AE (Eq. 4) is exactly a procedure of information
capability increasing. Thus, we take it as an indicator (see Figure 2(c)) in
turn to enhance $\beta$-VAE disentangling by adaptively adjusting the
objective of Eq. 2, where we change the hand-craft $C$ to a variable that can
automatically change in the optimization (Eq. 8). Such a mutually promoting
mechanism forms a closed-loop optimization flow.

###  4.2 Training of the Whole CL-Dis Framework

####  4.2.1 Phase-1: Diff-AE and $\beta$-VAE Pre-training

Following [64], the entire Diff-AE model is pre-trained on the various
datasets [40, 39, 96] to cover more semantic factors in the latent space of
$\mathbf{z_{sem}}$. But at this phase, $\mathbf{z_{sem}}$ is still far from
well-disentangled, which is just taken as a kind of coarse condition for
guiding diffusion generation [94].

Moreover, the co-pilot $\beta$-VAE is also pre-trained on the corresponding
datasets (i.e., if Diff-AE was trained on the facial data, $\beta$-VAE did the
same), following the objective of Eq. 1, its latent $\mathbf{z_{disen}}$ has
already enabled a preliminary disentangling capability, i.e., can distinguish
facial features of smile, azimuth, hair color, etc (Figure 2) to some degree.

####  4.2.2 Phase-2: Knowledge Distillation & Diffusion Feedback

To transfer the preliminary disentangling capability of $\mathbf{z_{disen}}$
into $\mathbf{z_{sem}}$, we use a distillation loss shown in Eq. 5 to smoothly
bridge Diff-AE and $\beta$-VAE. But this interaction is currently one-way,
which is still limited by the intrinsic shortcomings of $\beta$-VAE discussed
after Eq. 1 that it is hard to balance a trade-off and thus hurt the final
performance.

Therefore, we propose a novel diffusion feedback-driven loss
$\mathcal{L}_{fd}$ to further facilitate the disentanglement of the
representations ($\mathbf{z_{disen}}$, $\mathbf{z_{sem}}$) and break the
optimization bottleneck of the entire CL-Dis framework. Basically, inspired by
the work of understanding $\beta$-VAE [4] that modified the vanilla objective
by adding a hand-craft increasing controller $C$ (as shown in Eq. 2), we
propose to leverage the reverse diffusion of Diff-AE that adaptively increases
information capacity to build constraints to replace the original format.
Formally, we make the original $C$ auto-changeable, and control its value
dynamically according to the information capacity of the reverse diffusion
process as follows:

$C_{dyn}=f({E_{x_{t}}})=\begin{cases}C_{base}\cdot(\frac{E_{x_{0}}}{E_{x_{t}}}),&\text{
if }0<f({E_{x_{t}}})<C_{max}\\\ C_{max},&\text{ if }f({E_{x_{t}}})\geq
C_{max}\end{cases}$ (6)
$E_{{x_{0}}}=-\sum_{p_{{x_{0}}}\in\mathbb{R}}p_{{x_{0}}}logp_{{x_{0}}},\hskip
5.69054ptE_{{x_{t}}}=-\sum_{p_{{x_{t}}}\in\mathbb{R}}p_{{x_{t}}}logp_{{x_{t}}},$
(7)

where $E_{x_{0}}$ and $E_{x_{t}}$ denote the information entropy of the static
target image (i.e., Ground Truth) $x_{0}$ and the intermediate predicted
diffusion denoising image result $x_{t}$, respectively. The probability values
of $p_{x_{0}},p_{x_{t}}$ are calculated by flattening and probabilizing the
values of the image matrices across the dimension of $\mathbb{R}^{C\times
H\times W}$. Besides, $C_{base},C_{max}$ are the original empirical start-
point and the maximum upper-bound according to [4], and their influences are
studied in supplementary. We can intuitively see that with the denoising
diffusion processing, the entropy $E_{x_{t}}$ of the predicted image result
$x_{t}$ will gradually drop (i.e., information capacity increasing), and thus
the ratio of $\frac{E_{x_{0}}}{E_{x_{t}}}$ will increase to raise the
influence of $C_{dyn}$. So, the objective of $\beta$-VAE changes from Eq. 2 to
the following $\mathcal{L}_{fd}$,

$\displaystyle\mathcal{L}_{fd}(\theta,\phi,C_{dyn};\mathbf{x},\mathbf{z})=$
$\displaystyle\mathbb{E}_{q_{\phi}(\mathbf{z}\mid\mathbf{x})}\log
p_{\theta}(\mathbf{x}|\mathbf{z})-$
$\displaystyle\beta\big{|}D_{KL}\big{(}q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z})\big{)}-C_{dyn}\big{|},$
(8)

the above physical meaning is that: when KL-divergence drops to a certain
extent (i.e., disentangling ability encounters bottleneck in a naive VAE
latent), the dynamic controller $C_{dyn}$ will counteract this change by
increasing the value of the second item of Eq. 8. Intuitively, $C_{dyn}$ that
is controlled by diffusion feedback gradually highlights the penalty for KL-
divergence, thus improving the latent disentangling. Under this closed-loop
optimization flow (see Figure 2(c)), the gradually increasing information
capacity of diffusion adaptively enhances the disentanglement of latent
$\mathbf{z_{disen}}$, so as to further benefit the semantic code of
$\mathbf{z_{sem}}$. Please see supplementary for more theoretical analysis.

####  4.2.3 Phase-3: Semantics Directions Navigation

Figure 3: A self-supervised Navigation strategy is to discover the
interpretable shifts in the well-disentangled space of $\mathbf{z_{sem}}$. A
training sample consists of two images, produced by the Diff-AE with original
and shifted inputs. The images are given to a predictor $P$ that predicts the
direction index $k$ and shift magnitude $\delta$.

After getting the well-disentangled representations of $\mathbf{z_{sem}}$, the
next step is to identify its semantic meanings. Shifting a certain dimension
of the well-disentangled condition $\mathbf{z_{sem}}$ will result in a
generated result that only the corresponding property/factor changes. Thus, we
adopt an outer navigation branch to identify the traversal directions in a
global view. Specifically, we employ a learnable matrix $M$ (see Figure 3) to
manually append a shift on $\mathbf{z_{sem}}$. The shifted code
$\mathbf{z_{sem}^{\prime}}$, together with the original $\mathbf{z_{sem}}$, is
fed into our Diff-AE decoder to generate a pair of images, following the
reverse conditional diffusion process of $p(x_{t-1}|x_{t},\mathbf{z_{sem}})$.
A learnable predictor $P$ is then devised to predict the shift (i.e., semantic
direction) according to such generated paired images, with a reconstruction
loss. In this way, we clearly identify the inner disentangled semantics
directions in our CL-Dis for further fine-controllable generations.

###  4.3 New Disentanglement Metric

The conditional diffusion-based generation in our CL-Dis defines a clear
mapping between the disentangled representations and the generated image
results, so we can quantitatively measure the performance of disentanglement
with the generation difference. Following this idea, we quantify this
difference by using an optical flow tracking technique [76]: Given a well-
disentangled representation $\mathbf{z_{sem}^{a}}$ and its shift variant
$\mathbf{z_{sem}^{a^{\prime}}}$, we calculate the optical flow map $F$ between
their corresponding generated images. Then, we use a continuously changed
hyper-parameter as a dynamic threshold to distinguish these “ _changed
characteristics_ ” caused by shifted factors from the rest “ _frozen ones_ ”.
The average ratio of these two parts actually reflects the disentanglement
performance, the smaller the better. Due to the space limitations, more
operation details are in supplementary.

##  5 Experiment

###  5.1 Implementation Detail and Experiment Setting

For the backbone of Diff-AE, we take pre-trained weights in [64] as the
initial parameters with Adam optimizer and the learning rate of 1e-4. For the
co-pilot $\beta$-VAE, we adopt the same Adam optimizer and learning rate. For
the navigation branch, we refer to the GAN-based editing methods [79, 12] for
a fine-controllable generation. Differently, our CL-Dis has already in advance
got the well-disentangled latent $\mathbf{z_{sem}}$. The whole CL-Dis
framework is trained on 2 NVIDIA V100 GPUs with a batch size of 32. In
inference, the co-pilot $\beta$-VAE is discarded for efficiency. The dimension
of $\mathbf{z_{sem}}$ is set to 512 in our framework. To evaluate
disentanglement, we conduct experiments on both generation & perception tasks,
on synthetic & real data.

Datasets & Baselines. For the image generation task, we conduct experiments on
FFHQ [40], CelebA [39], LSUN Horse [96], LSUN Cars [96], Shape3D [42]. We
compare our proposed CL-Dis with VAE-based, GAN-based, and Diffusion-based
baselines, including $\beta$-VAE [30], FactorVAE [42], Beta-TCVAE [7], GAN
Edit v1 [79], GAN Edit v2 [12], InfoGAN-CR [54], Asyrp [48], GANspace (GS)
[27], NaviNeRF [88], DisCo [66], Diff-AE [64] and vector-based DisDiff [94].

For the perception task, we just choose the classic face recognition for
instantiation due to limited space and conduct experiments on CASIA-Webface
[95] and VGGFace2 [5] in comparison with the baseline of LFW [35].

Evaluation Metrics. To alleviate errors, we have multiple runs for each
method. We use 4 representative metrics of Factor-VAE score [42], DCI [19],
FID [29] and classification accuracy in face recognition, and our additionally
proposed new metric for validation. Note that our CL-Dis belongs to
dimension-/scalar-based DRL methods, so we just discuss but do not directly
compare with many vector-based ones [16, 77, 55, 93, 6] in the next experiment
sections.

###  5.2 Comparisons with SOTA DRL Methods

Table 1: Quantitative comparison results on CelebA dataset. Methods | FID
---|---
VAE [45] | 132.80
$\beta$-VAE [30] | 136.23
$\beta$-TCVAE [7] | 139.07
FactorVAE [42] | 134.52
CL-Dis (Ours) | 6.54

Table 2: Quantitative comparison results on FFHQ dataset. Methods | FID
---|---
NaviNeRF [88] | 13.00
CL-Dis (Ours) | 12.15

Quantitative Results. Table 2 and Table 2 report the results of Frechet
Inception Distance (FID) [29] to measure the quality of the generated images.
We see that our CL-Dis outperforms other typical disentangled methods in
generation quality, especially on the real-world face images of CelebA.

Besides, the quantitative comparison results of disentanglement under
different metrics are shown in Table 3. As shown in the table, CL-Dis
outperforms other competitors, demonstrating the method’s superior
disentanglement ability. Compared with the VAE-based methods, these methods
suffer from the the trade-off between generation and disentanglement [30, 4],
but CL-Dis does not have this problem due to its automatic optimization
mechanism. As for the GAN-based methods are typically limited by the latent
space of GAN [94]. The diffusion-based DisDiff is a vector-wise method, in
which representations $\\{z^{c}\\}$ could express more information than our
scalar-based ones $\mathbf{z_{sem}}$, but our CL-Dis still outperforms DisDiff
on Shape3D. And, our scalar-based disentangled representations are more light,
generic, robust, and easy to apply to different scenarios.

Table 3: Comparisons of disentanglement on the FactorVAE score and DCI disentanglement metrics (mean $\pm$ std, higher is better). We can see that our scalar-based CL-Dis even achieves comparable results with the co-current factor-based vector-wise DisDiff [94]. Method | Shapes3D |
---|---|---
FactorVAE score | DCI |
VAE-based: |  FactorVAE [42] | 0.840 ± 0.066 | 0.611 ± 0.082 |
$\beta$-TCVAE [7] | 0.873 ± 0.074 | 0.613 ± 0.114 |
GAN-based: |  InfoGAN-CR [54] | 0.587 ± 0.058 | 0.478 ± 0.055 |
Pre-trained GAN-based: |  GAN Edit v1 [79] | 0.805 ± 0.064 | 0.380 ± 0.062 |
GANspace [27] | 0.788 ± 0.091 | 0.284 ± 0.034 |
DisCo [66] | 0.877 ± 0.031 | 0.708 ± 0.048 |
Diffusion-based: |  DisDiff [94] | 0.902 ± 0.043 | 0.723 ± 0.013 |
CL-Dis (Ours) | 0.952 ± 0.017 | 0.731 ± 0.045 |
Figure 4: Fine-grained Disentanglement Results. The qualitative results
present the results of attribute manipulation, which demonstrate the semantic
manipulation on the FFHQ dataset. To clearly show the isolated and
“disentangled” effects, we use red circles to mark the corresponding changed
regions. Figure 5: Comparison of the “mouth”, “eyes” and “hair” attributes
with different methods on the FFHQ dataset. Since different methods pre-train
on different datasets, we thus compare the disentanglement with different
images according to the previous methods [79, 12, 48].

Qualitative Results. Unlike most previous papers that only test on the toy
synthetic datasets [30, 42, 7], our CL-Dis can be naturally applied to real
images. Figure 4 shows the effectiveness of our method on various datasets. To
clearly show the isolated and “disentangled” effects, we use red circles to
mark the corresponding changed regions. We can observe that CL-Dis can
disentangle and change the _isolated_ attributes continuously in an
unsupervised manner, such as _serious expression $\rightarrow$ smile, closing
eyes $\rightarrow$ opening eyes, w/o earrings $\rightarrow$ w/ earrings, etc._

Moreover, to validate the superiority of CL-Dis in the disentanglement
capability, Figure 5 also provides results for changing human facial
attributes to different characteristics along certain directions like “Mouth”,
“Eyes” and “hair”, in comparison with three baselines of GAN Edit v1 [79], GAN
Edit v2 [12], and Asyrp [48]. We can see that other methods like Asyrp are
prone to generate artifacts (bottom row of Figure 5), and their editing
results are not well-disentangled. For example, the whole face content has
changed for the schemes of GAN Edit v1 [79] and GAN Edit v2 [12]). In
contrast, our method consistently keeps the identity while just changing the
disentangled properties, achieving an “isolated and disentangled” effect.

To demonstrate the generalization ability of our CL-Dis, we also use different
datasets and different backbones (supplementary) for validation. As shown in
Figure 6, even for cars, horses, and buildings, CL-Dis still obtains a group
of fine-grained disentanglement results, where only the “color” attribute
changes across the traversal direction, proving the disentanglement
generalization capability.

Figure 6:  Generalization evaluation on extensive datasets. We visualize the
results of shifting the latent $\mathbf{z}_{sem}$ with our proposed method
along the common “Color” direction on LSUN Horse (upper row), LSUN Cars
(middle row), and CelebA (lower row) datasets, respectively. Figure 7:
Ablation study on the FFHQ dataset. We evaluate the effectiveness of different
components in our method by shifting different attributes, including “Hair”,
“Mouth” and “Eyes”. We can see that as more elements are added, both the
disentanglement performance and the generation quality improve: the first
column exists artifacts; the second and third columns change other contents
except shift attributes.

###  5.3 Ablation Study and More Analysis

Effectiveness of Different Elements. To analyze the effectiveness of the
proposed parts in our closed-loop disentanglement framework, we perform an
ablation to study the influence of Semantics Navigation, Knowledge
Distillation $\mathcal{L}_{dt}$, and Diffusion Feedback $\mathcal{L}_{fd}$. We
take FFHQ as the dataset to conduct these ablation studies.

As shown in Figure 7, due to the baseline of Diff-AE itself not having a
disentanglement ability, when shifting its latent we get generation results
with background artifacts. If we directly apply the semantics navigation [79]
to the baseline latent, the generation quality has been improved, but the
disentanglement is still not guaranteed. For example, when shifting along the
semantic direction of “Hair” (The 1st row of Figure 7), global features such
as face shape, skin texture, and hairstyle are variously entangled, resulting
in a younger appearance for the man. The last two schemes of adding knowledge
distillation and diffusion feedback have achieved better performance in
disentanglement, particularly for the final closed-loop scheme. We can see
that the final scheme together with $\mathcal{L}_{dt}$ & $\mathcal{L}_{fd}$
achieves a better disentanglement effect on the semantic direction of “Hair”,
compared to that only with $\mathcal{L}_{dt}$ which still changes face
characteristics in the traversal generations.

Analysis of Auto-Driven Feedback Mechanism. As introduced, the core
contribution of this work lies in an automatically disentangled learning
mechanism, achieved by a dynamic $C_{dyn}$ whose value depends on the
increasing information capacity in the reverse diffusion process. To prove
that, we visualize the changing curves of $C_{dyn}$ during the CL-Dis
training. As illustrated in Figure 8(a), the value of $C_{dyn}$ gradually
increases with the diffusion feedback, so as to highlight the penalty for KL-
divergence item in the co-pilot $\beta$-VAE optimization and thus improve
disentanglement performance of the entire framework.

###  5.4 New Metric $\&$ Discrimination Task Validation

Figure 8: (a) Visualization of the curve of changeable $C_{dyn}$ at different
iteration steps. $C_{dyn}$ gradually increases with the diffusion feedback to
improve the latent disentangling. (b) Averaged heatmaps of the pixel
differences between the original and the shifted images for the “eyes shift”
(left) and the “teeth shift” (right), revealing a good disentanglement of
ours.

New Metric Evaluation. As our work belongs to the track of unsupervised
representation disentanglement, not image generation or editing. The core of
this task is to fully understand the latent factors of the model to enhance
its fine-grained controllability. The qualitative results shown above have
illustrated that traversing along the disentangled representations results in
an “property isolated” generation effect. To measure such effect
quantitatively and show our method could progressively change the disentangled
properties while keeping others frozen, we propose a new metric (Sec. 4.3)
through the lens of optical flow [76]. As illustrated in Figure 9, the optical
flow visualization between the original and shifted image pairs on FFHQ
dataset [40], shows that our proposed CL-Dis method demonstrates effective
disentanglement, where shifting a single dimension of the latent space results
in isolated property changes in the generated output. In contrast, the
baseline approach (original Diff-AE) changes the entire background content of
the image.

Furthermore, we perform normalization on the optical flow map and utilize a
dynamic threshold to classify pixels as either exceeding or not exceeding it,
thereby obtaining the ratio. We make statistics based on the 100 sets, and
calculate the average ratio using a static threshold of 0.5, and get the
average value. The baseline quantitatively achieves 0.1606 in this new metric,
while our method is only 0.0177 (the lower the better), indicating that our
method has better disentanglement performance that only isolated property
changes when shifting across certain dimensions.

Figure 9: Evaluation on the proposed new disentanglement metric with the
‘hair’, ‘eyes’, and ‘smile’ attributes. From the optical flow tracking
visualization, we find traversing along our disentangled representations
results in an “property isolated” generation effect, while the baseline
exhibits a bad performance with a fully changed image content.

Locality of visual effects. To illustrate that the navigation along the
semantic directions results in an isolated and “disentangled” effect, for a
particular direction, we compute the per-pixel differences averaged over the
generation results with 20 shift magnitudes from a uniform grid in a range.
Figure 8(b) shows the averaged heatmap for the “Eyes” and “Mouth” directions.
We can see that our method achieves a good disentanglement, where shifting
“Eyes” just affects the eyes-related image content, and shifting “Mouth” just
affects the mouth-related image content.

Discrimination Task Validation. To validate whether the good disentanglement
achieved by our CL-Dis could benefit the discrimination task, we conducted
face recognition experiments on two large well-known datasets, CASIA-Webface
[95] and VGGFace2 [5]. We use the semantics encoder $E_{sem}$ of CL-Dis as an
off-the-shelf feature extractor and introduce two additional learnable linear
layers for fine-tuning. Following the same protocol used in [68], we perform a
fair comparison in Table 4 with mean classification accuracy as a metric. We
observe that the semantic encoder of our CL-Dis performs better than the
baseline of learning from scratch (LFS) on the widely adopted test benchmark
LFW [35], and even our $E_{sem}$ trained with only 5 epochs still outperforms
the LFS model with 50 epochs, which demonstrates that the well-disentangled
representations learned by CL-Dis indeed benefit discrimination tasks.

Table 4: Discrimination task (face recognition) validation. Our method with good disentangled representations achieves $>$2.0% gains even with fewer (5) training epochs compared to baseline. Dataset | CASIA-Webface (%) | VGGFace2 (%)
---|---|---
Baseline (50 epochs) | 89.45+-0.015 | 81.82+-0.019
Ours (5 epochs) | 91.72+-0.011 | 87.32+-0.014
Ours (50 epochs) | 94.00+-0.785 | 87.97+-0.012

##  6 Conclusion

In this paper, we present CL-Dis, an unsupervised closed-loop disentanglement
framework that achieves a strong DRL ability without relying on heavy
supervision or heuristic loss constraints. CL-Dis embeds a closed-loop
interconnection, where a pre-trained $\beta$-VAE empowers a preliminary
disentanglement ability to diffusion model with a distillation loss, while the
increasing information capability in the diffusion is taken as feedback, in
turn, to progressively improve $\beta$-VAE disentangling, forming a core
mutually-promoting mechanism. Besides, we also design a self-supervised
strategy to navigate semantic directions in the well-disentangled latent for a
controllable generation. A new metric is further introduced to measure the
disentanglement effect. Experimental results indicate that CL-Dis achieves a
state-of-the-art DRL performance, and its strong disentangled ability benefits
both generation and discrimination tasks, especially on real-world images.

## Supplementary Material

##  Appendix A Understanding the effects of our CL-Dis

Constrained optimization is important for enabling deep unsupervised models to
learn disentangled representations of the independent data generative factors
[30]:

(1) In the $\beta$-VAE framework, this corresponds to tuning the $\beta$
coefficient in Eq. 1 as follows,

$\displaystyle\mathcal{L}(\theta,\phi,\mathbf{x},\mathbf{z},\beta)=$
$\displaystyle\mathbb{E}_{q_{\phi}(\mathbf{z}|\mathbf{x})}\big{[}\log
p_{\theta}(\mathbf{x}|\mathbf{z})\big{]}-$ $\displaystyle\beta
D_{KL}\big{(}q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z})\big{)},$
(9)

$\beta$ can be taken as a mixing coefficient for balancing the magnitudes of
gradients from the reconstruction and the prior-matching components of the VAE
lower bound formulation during training. When $\beta$ is too low or too high
the model learns an entangled latent representation due to either too much or
too little capacity in the latent bottleneck [30].

(2) In the understanding $\beta$-VAE version with a controlled capacity
increase strategy [4], this constrained optimization corresponds to tuning the
$C$ coefficient in Eq 2 as follows,

$\displaystyle\mathcal{L}(\theta,\phi,C;\mathbf{x},\mathbf{z})=$
$\displaystyle\mathbb{E}_{q_{\phi}(\mathbf{z}\mid\mathbf{x})}\log
p_{\theta}(\mathbf{x}|\mathbf{z})-$
$\displaystyle\beta\big{|}D_{KL}\big{(}q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z})\big{)}-C\big{|},$
(10)

where the model is trained to reconstruct the images with samples from the
factor distribution, but with a range of different target encoding capacities
by pressuring the KL divergence to be at a controllable value, $C$. The
training objective combined maximizing the log-likelihood and minimizing the
absolute deviation from $C$ by linearly increasing it from a low value to a
high value over the course of training. The intuitive picture behind $C$ is
that it gradually adds more latent encoding capacity, enabling progressively
more factors of variation to be represented while retaining disentangling in
previously learned factors.

(3) In the proposed closed-loop unsupervised representation disentanglement
framework CL-Dis, this constrained optimization is totally automatic and
dynamic-adaptive. This is achieved under the mutual driving of reverse
diffusion of Diff-AE and VAE latent evolution. We make the constant $C$ in [4]
become changeable $C_{dyn}$, controlling its value dynamically according to
the information capacity variation in the reverse diffusion process as
follows,

$\displaystyle\mathcal{L}_{fd}(\theta,\phi,C_{dyn};\mathbf{x},\mathbf{z})=$
$\displaystyle\hskip
2.84526pt\mathbb{E}_{q_{\phi}(\mathbf{z}\mid\mathbf{x})}\log
p_{\theta}(\mathbf{x}|\mathbf{z})-$
$\displaystyle\beta\big{|}D_{KL}\big{(}q_{\phi}(\mathbf{z}|\mathbf{x})\|p_{\theta}(\mathbf{z})\big{)}-C_{dyn}\big{|},$
(11) $\displaystyle C_{dyn}\propto{I(x_{0};x_{t})}=$ $\displaystyle\sum\sum
p(x_{0},x_{t})\cdot\log\left(\frac{p(x_{0},x_{t})}{p(x_{0})\cdot
p(x_{t})}\right)$ (12) $\displaystyle s.t.\hskip
8.53581ptp_{\theta^{\prime}}(x_{t-1}|x_{t})=$ $\displaystyle\hskip
2.84526pt\mathcal{N}(x_{t};\mu_{\theta^{\prime}}(x_{t},t),\sigma_{t}I),$ (13)

where the $C_{dyn}$ is proportional to the information capability of the
reverse process in diffusion (here we use the mutual information
$I(x_{0};x_{t})$ to represent its value), which is parameterized by
${\theta^{\prime}}$ with a U-Net architecture and is a progressively
increasing process of information capacity (i.e., entropy decreasing).
Continuing this intuitive picture, we can imagine that with the capacity of
the information capability gradually increasing in the reverse diffusion, the
dynamic $C_{dyn}$ would continue to utilize those extra indicators for an
increasingly precise encoding, where a larger improvement can be obtained by
disentangling different factors of variations in the latent.

##  Appendix B Implementation Details of the New Metric

To measure the disentanglement effect quantitatively, especially under the
unsupervised setting where no factor annotations are accessible, we propose a
new label-free metric with the help of optical flow [76]. The high-level idea
is that we leverage optical flow to track the obviously and continually
changed content to evaluate whether or not these changes are “property
isolated”. First, we calculate the optical flow between the original image and
the shifted image. Then, we normalize the values of the optical flow map.
After that, we apply a threshold to the optical flow. Points in the flow
exceeding this threshold are considered to have a large shift magnitude (i.e.,
those disentangled attribute points), while those below it are considered to
have a small shift magnitude (i.e., those rest frozen attribute points). We
calculate the ratio of the number of flow points exceeding this threshold to
those below it and use this ratio value as our new metric score. The lower the
better, and which exactly corresponds to the definition of disentangled
representations where a change in a single unit of the representation
corresponds to a change in a single factor of variation of the data while
being invariant to others. The specific algorithm is shown in Algorithm 1 and
Algorithm 2. Empirically and experimentally, we set the threshold as 0.5 for
the calculation of the proposed new metric.

Algorithm 1 Compute New Metric Based on Optical Flow

1:function ComputeNewMetric($image$, $shifted\\_image$, $threshold$)

2: $flow\leftarrow\textsc{OpticalFlow}(image,shifted\\_image)$

3: $normalized\\_flow\leftarrow\textsc{Norm}(flow)$

4: $large\\_shift\\_count\ \mathcal{L}\leftarrow 0$

5: $small\\_shift\\_count\ \mathcal{S}\leftarrow 0$

6: for all $flow\\_point$ in $normalized\\_flow$ do

7: if $\textsc{Magnitude}(flow\\_point)>threshold$ then

8: $\mathcal{L}\leftarrow\mathcal{L}+1$

9: else

10: $\mathcal{S}\leftarrow\mathcal{S}+1$

11: end if

12: end for

13: $metric\hskip 2.84526ptscore\leftarrow\ \mathcal{L}/\mathcal{S}$

14: return $metric\hskip 2.84526ptscore$

15:end function

Algorithm 2 Optical Flow Normalization

1:function Norm($flow$)

2: $max\\_norm\leftarrow\textsc{Max}(\sqrt{\sum{flow^{2}}})$

3: $normalized\\_flow\leftarrow flow/max\\_norm$

4: return $normalized\\_flow$

5:end function

##  Appendix C More Experimental Results

###  C.1 Influence of $C_{base}$ and $C_{max}$ in Our CL-Dis

Table 5: Influence study with different values of $C_{base}$ and $C_{max}$ in our CL-Dis framework on the Cars3D dataset. We evaluate the disentanglement performance of different schemes using the metrics of FactorVAE score and DCI (mean $\pm$ std, higher is better). $C_{base}$ | $C_{max}$ | FactorVAE score | DCI
---|---|---|---
2 | 5 | 0.689 ± 0.038 | 0.094 ± 0.009
5 | 15 | 0.745 ± 0.042 | 0.175 ± 0.012
10 | 11 | 0.762 ± 0.068 | 0.179 ± 0.014
10 | 12 | 0.781 ± 0.035 | 0.182 ± 0.008
10 | 25 | 0.794 ± 0.040 | 0.186 ± 0.007
20 | 50 | 0.734 ± 0.024 | 0.163 ± 0.015
Figure 10: Generalization evaluation on extensive datasets. We visualize the results of shifting the disentangled latent $\mathbf{z_{sem}}$ with our proposed method. From top to bottom, we illustrate the “property isolated” effects along glasses, mouth, hair, car color, hair color, and face illumination, respectively. Table 6: Comparisons of disentanglement on the FactorVAE score and DCI disentanglement metrics (mean $\pm$ std, higher is better). We can see that our scalar-based CL-Dis even achieves comparable results with the co-current factor-based vector-wise DisDiff [94]. Method | Shapes3D |
---|---|---
FactorVAE score | DCI |
VAE-based: |  FactorVAE [42] | 0.840 ± 0.066 | 0.611 ± 0.082 |
$\beta$-TCVAE [7] | 0.873 ± 0.074 | 0.613 ± 0.114 |
GAN-based: |  InfoGAN-CR [54] | 0.587 ± 0.058 | 0.478 ± 0.055 |
Pre-trained GAN-based: |  GAN Edit v1 [79] | 0.805 ± 0.064 | 0.380 ± 0.062 |
GANspace [27] | 0.788 ± 0.091 | 0.284 ± 0.034 |
DisCo [66] | 0.877 ± 0.031 | 0.708 ± 0.048 |
Diffusion-based: |  DisDiff [94] | 0.902 ± 0.043 | 0.723 ± 0.013 |
CL-Dis (Ours) | 0.952 ± 0.017 | 0.731 ± 0.045 |

As we discussed in the main paper, we propose a novel diffusion feedback-
driven loss $\mathcal{L}_{fd}$ to further facilitate the disentanglement of
the representations. Formally, we make the original $C$ in $\beta$-VAE auto-
changeable, and control its value dynamically according to the information
capacity of the reverse diffusion process as follows:

$C_{dyn}=f({E_{x_{t}}})=\begin{cases}C_{base}\cdot(\frac{E_{x_{0}}}{E_{x_{t}}}),&\text{
if }0<f({E_{x_{t}}})<C_{max}\\\ C_{max},&\text{ if }f({E_{x_{t}}})\geq
C_{max}\end{cases}$ (14)
$E_{{x_{0}}}=-\sum_{p_{{x_{0}}}\in\mathbb{R}}p_{{x_{0}}}logp_{{x_{0}}},\hskip
5.69054ptE_{{x_{t}}}=-\sum_{p_{{x_{t}}}\in\mathbb{R}}p_{{x_{t}}}logp_{{x_{t}}},$
(15)

where $C_{base}$, $C_{max}$ are the original empirical start-point and the
maximum upper-bound according to [4], and their influences are evaluated as
follows. The results depicted in Table 5 demonstrate that our method is robust
enough to achieve a stable satisfactory performance with different values of
$C_{base}$, $C_{max}$. Besides, these experimental results reveal that setting
$C_{base}$ to 10 consistently yields optimal disentanglement outcomes.
Moreover, our assessment also extends to examining the impact of the upper-
bound value, $C_{max}$. As shown in Table 5, reducing the gap between
$C_{max}$ and $C_{base}$ hampers the overall disentanglement performance
(0.794 vs. 0.762 in FactorVAE score). We analyze that, the variable scope of
the dynamic $C_{dyn}$ will be limited due to the relatively small range
between $C_{max}$ and $C_{base}$. This limitation may negatively impact the
control of KL-divergence, consequently hampering the overall model’s ability
to effectively disentangle the inner latent space more.

###  C.2 More Subjective Results

In this section, we further provide more visualization results to demonstrate
the generalization ability of our CL-Dis on different datasets, including FFHQ
[40], LSUN Cars [96] and CelebA [39]. As shown in Figure 10, our CL-Dis could
generate fine-grained disentangled results along the “glasses, mouth, hair,
car color, hair color, and face illumination” attribute changes across the
traversal direction, which fully illustrates the well-generalized
disentanglement capability of our method.

###  C.3 Dimension Ablation on Shape3D

We further conduct Dimension Ablation on Shapes3D dataset in Table 7. We can
see below that reducing the dimension of $z_{sem}$ leads to minor changes,
which indicates our method is robust enough.

Dimension | FactorVAE score | DCI
---|---|---
CL-Dis (512) | 0.840 ± 0.066 | 0.611 ± 0.082
CL-Dis (100) | 0.842 ± 0.042 | 0.547 ± 0.051
CL-Dis (32) | 0.837 ± 0.053 | 0.582 ± 0.049
Table 7: Ablation study for the dimension of $z_{sem}$ on the Shapes3D
dataset.

##  Appendix D Limitation Discussions

In Table 6 (corresponds to Table 2 in the main manuscript), we observe our
method encounters an unsatisfactory performance on the MPI3D dataset, which
indicates that our method may be distracted by the real-world natural
environment scenarios. We discuss the potential reason that the collected data
in MPI3D [23], focus on bridging the gap between simulation and reality, which
distinguishes it apart from various other synthetic toy datasets [65, 42]. In
detail, real-world recordings in MPI3D introduce more extra environment
interruptions, such as chromatic aberrations in cameras and complex surface
properties of objects (e.g., reflectance, radiance, and irradiance). These
interruptions pose challenges in learning disentangled representations due to
the mixing of factors and distortions during the real-world recordings,
especially for scalar-based methods like ours [23]. This unsolved limitation
of our method will be left as future research work.

## References

  * Ahn et al. [2019] Sungsoo Ahn, Shell Xu Hu, Andreas Damianou, Neil D Lawrence, and Zhenwen Dai.  Variational information distillation for knowledge transfer.  In _CVPR_ , pages 9163–9171, 2019.
  * Bao et al. [2022] Zhipeng Bao, Martial Hebert, and Yu-Xiong Wang.  Generative modeling for multi-task visual learning.  In _ICML_ , pages 1537–1554. PMLR, 2022.
  * Bengio et al. [2013] Yoshua Bengio, Aaron Courville, and Pascal Vincent.  Representation learning: A review and new perspectives.  _TPAMI_ , 35(8):1798–1828, 2013.
  * Burgess et al. [2017] Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner.  Understanding disentangling in $\beta$-vae.  _NeurIPS Workshop on Learning Disentangled Representations_ , 2017.
  * Cao et al. [2018] Qiong Cao, Li Shen, Weidi Xie, Omkar M. Parkhi, and Andrew Zisserman.  VGGFace2: A dataset for recognising faces across pose and age.  In _International Conference on Automatic Face and Gesture Recognition_ , 2018.
  * Chen et al. [2023] Hong Chen, Yipeng Zhang, Xin Wang, Xuguang Duan, Yuwei Zhou, and Wenwu Zhu.  Disenbooth: Disentangled parameter-efficient tuning for subject-driven text-to-image generation.  _arXiv preprint arXiv:2305.03374_ , 2023.
  * Chen et al. [2018] Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud.  Isolating sources of disentanglement in variational autoencoders.  _NeurIPS_ , 31, 2018.
  * Chen et al. [2016] Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel.  Infogan: Interpretable representation learning by information maximizing generative adversarial nets.  _NeurIPS_ , 29, 2016.
  * Chen et al. [2022] Zikun Chen, Ruowei Jiang, Brendan Duke, Han Zhao, and Parham Aarabi.  Exploring gradient-based multi-directional controls in gans.  In _ECCV_ , pages 104–119. Springer, 2022.
  * Cheng et al. [2020] Pengyu Cheng, Martin Renqiang Min, Dinghan Shen, Christopher Malon, Yizhe Zhang, Yitong Li, and Lawrence Carin.  Improving disentangled text representation learning with information-theoretic guidance.  _arXiv preprint arXiv:2006.00693_ , 2020.
  * Cheng et al. [2017] Yu Cheng, Duo Wang, Pan Zhou, and Tao Zhang.  A survey of model compression and acceleration for deep neural networks.  _arXiv preprint arXiv:1710.09282_ , 2017.
  * Cherepkov et al. [2021] Anton Cherepkov, Andrey Voynov, and Artem Babenko.  Navigating the gan parameter space for semantic image editing.  In _CVPR_ , pages 3671–3680, 2021.
  * Dai et al. [2023] Josef Dai, Xuehai Pan, Ruiyang Sun, Jiaming Ji, Xinbo Xu, Mickel Liu, Yizhou Wang, and Yaodong Yang.  Safe rlhf: Safe reinforcement learning from human feedback.  _arXiv preprint arXiv:2310.12773_ , 2023.
  * Dasgupta and Sabato [2020] Sanjoy Dasgupta and Sivan Sabato.  Robust learning from discriminative feature feedback.  In _International Conference on Artificial Intelligence and Statistics_ , pages 973–982. PMLR, 2020.
  * Dasgupta et al. [2018] Sanjoy Dasgupta, Akansha Dey, Nicholas Roberts, and Sivan Sabato.  Learning from discriminative feature feedback.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Denton et al. [2017] Emily L Denton et al.  Unsupervised learning of disentangled representations from video.  _NeurIPS_ , 30, 2017.
  * Driess et al. [2023] Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al.  Palm-e: An embodied multimodal language model.  _arXiv preprint arXiv:2303.03378_ , 2023.
  * Dupont [2018] Emilien Dupont.  Learning disentangled joint continuous and discrete representations.  _NeurIPS_ , 31, 2018.
  * Eastwood and Williams [2018] Cian Eastwood and Christopher KI Williams.  A framework for the quantitative evaluation of disentangled representations.  In _International conference on learning representations_ , 2018.
  * Egan [1989] Kieran Egan.  Memory, imagination, and learning: Connected by the story.  _Phi Delta Kappan_ , 70(6):455–459, 1989.
  * Fleming [2014] Roland W Fleming.  Visual perception of materials and their properties.  _Vision research_ , 94:62–75, 2014.
  * Goetschalckx et al. [2019] Lore Goetschalckx, Alex Andonian, Aude Oliva, and Phillip Isola.  Ganalyze: Toward visual definitions of cognitive image properties.  In _ICCV_ , pages 5744–5753, 2019.
  * Gondal et al. [2019] Muhammad Waleed Gondal, Manuel Wuthrich, Djordje Miladinovic, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer.  On the transfer of inductive bias from simulation to the real world: a new disentanglement dataset.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Gonzalez-Garcia et al. [2018] Abel Gonzalez-Garcia, Joost Van De Weijer, and Yoshua Bengio.  Image-to-image translation for cross-domain disentanglement.  _NeurIPS_ , 31, 2018.
  * Goodfellow et al. [2014] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  _NeurIPS_ , 27, 2014.
  * Gou et al. [2021] Jianping Gou, Baosheng Yu, Stephen J Maybank, and Dacheng Tao.  Knowledge distillation: A survey.  _IJCV_ , 129:1789–1819, 2021.
  * Härkönen et al. [2020] Erik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, and Sylvain Paris.  Ganspace: Discovering interpretable gan controls.  _NeurIPS_ , 33:9841–9850, 2020.
  * He et al. [2017] Ruidan He, Wee Sun Lee, Hwee Tou Ng, and Daniel Dahlmeier.  An unsupervised neural attention model for aspect extraction.  In _ACL_ , pages 388–397, 2017.
  * Heusel et al. [2017] 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.
  * Higgins et al. [2016] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.  beta-vae: Learning basic visual concepts with a constrained variational framework.  In _ICLR_ , 2016.
  * Higgins et al. [2018] Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerchner.  Towards a definition of disentangled representations.  _arXiv preprint arXiv:1812.02230_ , 2018.
  * Ho and Salimans [2022] Jonathan Ho and Tim Salimans.  Classifier-free diffusion guidance.  _arXiv preprint arXiv:2207.12598_ , 2022.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _NeurIPS_ , 33:6840–6851, 2020.
  * Hsu and Glass [2018] Wei-Ning Hsu and James Glass.  Disentangling by partitioning: A representation learning framework for multimodal sensory data.  _arXiv preprint arXiv:1805.11264_ , 2018.
  * Huang et al. [2008] Gary B Huang, Marwan Mattar, Tamara Berg, and Eric Learned-Miller.  Labeled faces in the wild: A database forstudying face recognition in unconstrained environments.  In _Workshop on faces in’Real-Life’Images: detection, alignment, and recognition_ , 2008.
  * Jeon et al. [2021] Insu Jeon, Wonkwang Lee, Myeongjang Pyeon, and Gunhee Kim.  Ib-gan: Disentangled representation learning with information bottleneck generative adversarial networks.  In _AAAI_ , pages 7926–7934, 2021.
  * Ji et al. [2023] Jiaming Ji, Tianyi Qiu, Boyuan Chen, Borong Zhang, Hantao Lou, Kaile Wang, Yawen Duan, Zhonghao He, Jiayi Zhou, Zhaowei Zhang, et al.  Ai alignment: A comprehensive survey.  _arXiv preprint arXiv:2310.19852_ , 2023.
  * Ji et al. [2021] Mingi Ji, Byeongho Heo, and Sungrae Park.  Show, attend and distill: Knowledge distillation via attention-based feature matching.  In _AAAI_ , pages 7945–7952, 2021.
  * Karras et al. [2017] 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.
  * Karras et al. [2019] Tero Karras, Samuli Laine, and Timo Aila.  A style-based generator architecture for generative adversarial networks.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 4401–4410, 2019.
  * Katakkar et al. [2021] Anurag Katakkar, Clay H Yoo, Weiqin Wang, Zachary C Lipton, and Divyansh Kaushik.  Practical benefits of feature feedback under distribution shift.  _arXiv preprint arXiv:2110.07566_ , 2021.
  * Kim and Mnih [2018] Hyunjik Kim and Andriy Mnih.  Disentangling by factorising.  In _ICML_ , pages 2649–2658. PMLR, 2018.
  * Kim et al. [2019] Minyoung Kim, Yuting Wang, Pritish Sahu, and Vladimir Pavlovic.  Relevance factor vae: Learning and identifying disentangled factors.  _arXiv preprint arXiv:1902.01568_ , 2019.
  * Kim et al. [2022] Yeongmin Kim, Dongjun Kim, HyeonMin Lee, and Il-chul Moon.  Unsupervised controllable generation with score-based diffusion models: Disentangled latent code guidance.  In _NeurIPS 2022 Workshop on Score-Based Methods_ , 2022.
  * Kingma and Welling [2013] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Knox and Stone [2011] W Bradley Knox and Peter Stone.  Augmenting reinforcement learning with human feedback.  In _ICML 2011 Workshop on New Developments in Imitation Learning (July 2011)_ , page 3, 2011.
  * Kumar et al. [2018] Abhishek Kumar, Prasanna Sattigeri, and Avinash Balakrishnan.  Variational inference of disentangled latent concepts from unlabeled observations.  _ICLR_ , 2018.
  * Kwon et al. [2022] Mingi Kwon, Jaeseok Jeong, and Youngjung Uh.  Diffusion models already have a semantic latent space.  _arXiv preprint arXiv:2210.10960_ , 2022.
  * Lake et al. [2017] Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel J Gershman.  Building machines that learn and think like people.  _Behavioral and brain sciences_ , 40:e253, 2017.
  * Lee et al. [2018] Hsin-Ying Lee, Hung-Yu Tseng, Jia-Bin Huang, Maneesh Singh, and Ming-Hsuan Yang.  Diverse image-to-image translation via disentangled representations.  In _ECCV_ , pages 35–51, 2018.
  * Lee and Sewell [2023] Ping-Shien Lee and David K Sewell.  A revised diffusion model for conflict tasks.  _Psychonomic Bulletin & Review_, pages 1–31, 2023.
  * Lee et al. [2020] Wonkwang Lee, Donggyun Kim, Seunghoon Hong, and Honglak Lee.  High-fidelity synthesis with disentangled representation.  In _ECCV_ , pages 157–174. Springer, 2020.
  * Li et al. [2021] Minjun Li, Yanghua Jin, and Huachun Zhu.  Surrogate gradient field for latent space manipulation.  In _CVPR_ , pages 6529–6538, 2021.
  * Lin et al. [2019] Zinan Lin, Kiran Koshy Thekumparampil, Giulia Fanti, and Sewoong Oh.  Infogan-cr: Disentangling generative adversarial networks with contrastive regularizers.  _arXiv preprint arXiv:1906.06034_ , page 60, 2019.
  * Liu et al. [2021a] Liu Liu, Jiangtong Li, Li Niu, Ruicong Xu, and Liqing Zhang.  Activity image-to-video retrieval by disentangling appearance and motion.  In _AAAI_ , pages 2145–2153, 2021a.
  * Liu et al. [2021b] Yahui Liu, Enver Sangineto, Yajing Chen, Linchao Bao, Haoxian Zhang, Nicu Sebe, Bruno Lepri, Wei Wang, and Marco De Nadai.  Smoothing the disentangled latent style space for unsupervised image-to-image translation.  In _CVPR_ , pages 10785–10794, 2021b.
  * Locatello et al. [2019] Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem.  Challenging common assumptions in the unsupervised learning of disentangled representations.  In _international conference on machine learning_ , pages 4114–4124. PMLR, 2019.
  * Luo et al. [2023] Zhipeng Luo, Yazhou He, Yanbing Xue, Hongjun Wang, Milos Hauskrecht, and Tianrui Li.  Hierarchical active learning with qualitative feedback on regions.  _IEEE Transactions on Human-Machine Systems_ , 2023.
  * Materzyńska et al. [2022] Joanna Materzyńska, Antonio Torralba, and David Bau.  Disentangling visual and written concepts in clip.  In _CVPR_ , pages 16410–16419, 2022.
  * Pearson [2019] Joel Pearson.  The human imagination: the cognitive neuroscience of visual mental imagery.  _Nature reviews neuroscience_ , 20(10):624–634, 2019.
  * Pierobon and Akyildiz [2012] Massimiliano Pierobon and Ian F Akyildiz.  Capacity of a diffusion-based molecular communication system with channel memory and molecular noise.  _IEEE Transactions on Information Theory_ , 59(2):942–954, 2012.
  * Plumerault et al. [2020] Antoine Plumerault, Hervé Le Borgne, and Céline Hudelot.  Controlling generative models with continuous factors of variations.  _ICLR_ , 2020.
  * Poulis and Dasgupta [2017] Stefanos Poulis and Sanjoy Dasgupta.  Learning with feature feedback: from theory to practice.  In _Artificial Intelligence and Statistics_ , pages 1104–1113. PMLR, 2017.
  * Preechakul et al. [2022] Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn.  Diffusion autoencoders: Toward a meaningful and decodable representation.  In _CVPR_ , pages 10619–10629, 2022.
  * Reed et al. [2015] Scott E Reed, Yi Zhang, Yuting Zhang, and Honglak Lee.  Deep visual analogy-making.  _Advances in neural information processing systems_ , 28, 2015.
  * Ren et al. [2021] Xuanchi Ren, Tao Yang, Yuwang Wang, and Wenjun Zeng.  Learning disentangled representation by exploiting pretrained generative models: A contrastive learning view.  _ICLR_ , 2021.
  * Savva et al. [2019] Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, et al.  Habitat: A platform for embodied ai research.  In _ICCV_ , pages 9339–9347, 2019.
  * Schroff et al. [2015] Florian Schroff, Dmitry Kalenichenko, and James Philbin.  Facenet: A unified embedding for face recognition and clustering.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 815–823, 2015.
  * Shao et al. [2020] Huajie Shao, Shuochao Yao, Dachun Sun, Aston Zhang, Shengzhong Liu, Dongxin Liu, Jun Wang, and Tarek Abdelzaher.  Controlvae: Controllable variational autoencoder.  In _International Conference on Machine Learning_ , pages 8655–8664. PMLR, 2020.
  * Shen and Zhou [2021] Yujun Shen and Bolei Zhou.  Closed-form factorization of latent semantics in gans.  In _CVPR_ , pages 1532–1540, 2021.
  * Shen et al. [2020] Yujun Shen, Jinjin Gu, Xiaoou Tang, and Bolei Zhou.  Interpreting the latent space of gans for semantic face editing.  In _CVPR_ , pages 9243–9252, 2020.
  * Shi et al. [2022] Yichun Shi, Xiao Yang, Yangyue Wan, and Xiaohui Shen.  Semanticstylegan: Learning compositional generative priors for controllable image synthesis and editing.  In _CVPR_ , pages 11254–11264, 2022.
  * Song et al. [2020] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  In _ICLR_ , 2020.
  * Song et al. [2023a] Yue Song, Andy Keller, Nicu Sebe, and Max Welling.  Latent traversals in generative models as potential flows.  _ICML_ , 2023a.
  * Song et al. [2023b] Yue Song, Jichao Zhang, Nicu Sebe, and Wei Wang.  Householder projector for unsupervised latent semantics discovery.  In _ICCV_ , pages 7712–7722, 2023b.
  * Teed and Deng [2020] Zachary Teed and Jia Deng.  Raft: Recurrent all-pairs field transforms for optical flow.  In _ECCV_ , pages 402–419. Springer, 2020.
  * Tran et al. [2017] Luan Tran, Xi Yin, and Xiaoming Liu.  Disentangled representation learning gan for pose-invariant face recognition.  In _CVPR_ , pages 1415–1424, 2017.
  * Tsai et al. [2018] Yao-Hung Hubert Tsai, Paul Pu Liang, Amir Zadeh, Louis-Philippe Morency, and Ruslan Salakhutdinov.  Learning factorized multimodal representations.  _ICLR_ , 2018.
  * Voynov and Babenko [2020] Andrey Voynov and Artem Babenko.  Unsupervised discovery of interpretable directions in the gan latent space.  In _ICML_ , pages 9786–9796. PMLR, 2020.
  * Wang et al. [2018] Hui Wang, Hanbin Zhao, Xi Li, and Xu Tan.  Progressive blockwise knowledge distillation for neural network acceleration.  In _IJCAI_ , pages 2769–2775, 2018.
  * Wang et al. [2022a] Tengfei Wang, Yong Zhang, Yanbo Fan, Jue Wang, and Qifeng Chen.  High-fidelity gan inversion for image attribute editing.  In _CVPR_ , pages 11379–11388, 2022a.
  * Wang et al. [2022b] Xin Wang, Hong Chen, Si’ao Tang, Zihao Wu, and Wenwu Zhu.  Disentangled representation learning.  _arXiv preprint arXiv:2211.11695_ , 2022b.
  * Wang et al. [2023] Zhizhong Wang, Lei Zhao, and Wei Xing.  Stylediffusion: Controllable disentangled style transfer via diffusion models.  In _ICCV_ , pages 7677–7689, 2023.
  * Wu et al. [2020] Jiawei Wu, Xiaoya Li, Xiang Ao, Yuxian Meng, Fei Wu, and Jiwei Li.  Improving robustness and generality of nlp models using disentangled representations.  _arXiv preprint arXiv:2009.09587_ , 2020.
  * Wu et al. [2023a] Qiucheng Wu, Yujian Liu, Handong Zhao, Ajinkya Kale, Trung Bui, Tong Yu, Zhe Lin, Yang Zhang, and Shiyu Chang.  Uncovering the disentanglement capability in text-to-image diffusion models.  In _CVPR_ , pages 1900–1910, 2023a.
  * Wu et al. [2023b] Yankun Wu, Yuta Nakashima, and Noa Garcia.  Not only generative art: Stable diffusion for content-style disentanglement in art analysis.  In _Proceedings of the 2023 ACM International Conference on Multimedia Retrieval_ , pages 199–208, 2023b.
  * Wu et al. [2021] Zongze Wu, Dani Lischinski, and Eli Shechtman.  Stylespace analysis: Disentangled controls for stylegan image generation.  In _CVPR_ , pages 12863–12872, 2021.
  * Xie et al. [2023] Baao Xie, Bohan Li, Zequn Zhang, Junting Dong, Xin Jin, Jingyu Yang, and Wenjun Zeng.  Navinerf: Nerf-based 3d representation disentanglement by latent semantic navigation.  _arXiv preprint arXiv:2304.11342_ , 2023.
  * Xu et al. [2023] Xingqian Xu, Zhangyang Wang, Gong Zhang, Kai Wang, and Humphrey Shi.  Versatile diffusion: Text, images and variations all in one diffusion model.  In _ICCV_ , pages 7754–7765, 2023.
  * Xu et al. [2022] Zipeng Xu, Tianwei Lin, Hao Tang, Fu Li, Dongliang He, Nicu Sebe, Radu Timofte, Luc Van Gool, and Errui Ding.  Predict, prevent, and evaluate: Disentangled text-driven image manipulation empowered by pre-trained vision-language model.  In _CVPR_ , pages 18229–18238, 2022.
  * Yan et al. [2016] Xinchen Yan, Jimei Yang, Kihyuk Sohn, and Honglak Lee.  Attribute2image: Conditional image generation from visual attributes.  In _ECCV_ , pages 776–791. Springer, 2016.
  * Yang et al. [2021] Tao Yang, Xuanchi Ren, Yuwang Wang, Wenjun Zeng, and Nanning Zheng.  Towards building a group-based unsupervised representation disentanglement framework.  _ICLR_ , 2021.
  * Yang et al. [2023a] Tao Yang, Yuwang Wang, Cuiling Lan, Yan Lu, and Nanning Zheng.  Vector-based representation is the key: A study on disentanglement and compositional generalization.  _arXiv preprint arXiv:2305.18063_ , 2023a.
  * Yang et al. [2023b] Tao Yang, Yuwang Wang, Yan Lv, and Nanning Zh.  Disdiff: Unsupervised disentanglement of diffusion probabilistic models.  _arXiv preprint arXiv:2301.13721_ , 2023b.
  * Yi et al. [2014] Dong Yi, Zhen Lei, Shengcai Liao, and Stan Z Li.  Learning face representation from scratch.  _arXiv preprint arXiv:1411.7923_ , 2014.
  * Yu et al. [2015] Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao.  Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop.  _arXiv preprint arXiv:1506.03365_ , 2015.
  * Zhang et al. [2022] Yuhao Zhang, Ying Zhang, Wenya Guo, Xiangrui Cai, and Xiaojie Yuan.  Learning disentangled representation for multimodal cross-domain sentiment analysis.  _TNNLS_ , 2022.
  * Zhu et al. [2021] Xinqi Zhu, Chang Xu, and Dacheng Tao.  Where and what? examining interpretable disentangled representations.  In _CVPR_ , pages 5861–5870, 2021.


#  Performance best practices using Java AWS lambda
Measuring enhancements of different optimization techniques combining java
with AWS lambda

Juan Mera Menéndez, Martin Bartlett

###### Abstract

Despite its already widespread popularity, it continues to gain adoption. More
and more developers and architects continue to adopt and apply the FaaS
(Function as a Service) model in cloud solutions. The most extensively used
FaaS service is AWS Lambda, provided by Amazon Web Services. Moreover, despite
the new trends in programming languages, Java still maintains a significant
share of usage. The main problem that arises when using these two technologies
together is widely known: significant latencies and the dreaded cold start.
However, it is possible to greatly mitigate this problem without dedicating
too much effort. In this article, various techniques, strategies and
approaches will be studied with the aim of reducing the cold start and
significantly improving the performance of Lambda functions with Java.
Starting from a system that involves AWS lambda, java, DynamoDB and Api
Gateway. Each approach will be tested independently, analyzing its impact
through load tests. Subsequently, they will be tested in combination in an
effort to achieve the greatest possible performance improvement.

###### Index Terms:

AWS Lambda, AWS, Amazon Web Services, Java, performance, cold start, FaaS,
Function as a Service, Serverless, GraalVM, Snapstart, JAVA_TOOL_OPTIONS

##  I Introduction

Following the trend of recent years, the serverless approach [9] continues to
gain popularity due to its scalability, flexibility, agility, and other well-
known benefits. According to the survey conducted by Datadog in August 2023
[13], one of the most popularly used services within this architectural model
is AWS Lambda (provided by Amazon Web Services) and over a 10% of Lambda
functions are written in Java. [11] Furthermore Java is the primary language
for 30% of professional developers.
Being the performance one of the main criteria when comparing solutions, the
survey shows that cold start of java lambdas are roughly twice as long as the
cold starts of other languages like Node.js and Python [3, 4](which are the
most used languages with AWS lambda). Indeed, it becomes evident why
performance is the first aspect brought to the table when discussing the
combination of Java with AWS Lambda. And rightfully so, if it’s not taken into
account when building the system, it’s common to encounter time-out issues in
API Gateway due to excessively high latency of the functions. Furthermore,
another noteworthy benefit of improving the performance of Lambda functions is
the reduction in their cost. This is because the cost model for this service
is based on computing time. By reducing this time, the cost is also reduced.
Thus, the primary goal of the study is to optimize the way Java code is
executed on AWS Lambda and provide a set of best practices that assist this
approach in being competitive in terms of performance. To achieve that goal, a
series of techniques, approaches or strategies are proposed with the aim of
improving function performance and mitigating cold starts:

  * •

Profile the function to see which memory configuration is better suited in
terms of performance.

  * •

Use Snapstart to reduce the cold-start latency.

  * •

Choose the Arm64 Lambda architecture against the default x86_64 architecture.

  * •

Use the AWS SDK v2 for Java.

  * •

Pre-compile the source code with GraalVM to avoid initialize all classes in
runtime.

  * •

Java Lambda function customization settings: JAVA_TOOL_OPTIONS environment
variable.

  * •

Other good practices

The improvement caused by each technique will be measured independently since
some of them are not compatible with each other. The advantages and drawbacks
of the techniques will be discussed as well. For the implementation of each
technique or approach, we start with an initial, unoptimized system, which
will be explained in the following section. The performance of this system has
also been measured and will be used to calculate the improvements experienced
relative to it. After this, the techniques will be combined, and the result
will be measured.

Figure 1: Initial application architecture

##  II Initial application

A common use case for AWS Lambda is to integrate with DynamoDB, allowing
functions to perform operations on data stored in DynamoDB. Following this
approach, an initial solution has been built to measure its initial
performance and demonstrate the performance improvements generated as a result
of the proposed best practices.
The mentioned system consists of a simple e-commerce solution, compound of
three services with functionality which translates in one function for each
service. All the services are self-contained and individual, offering an API
through of which the services send messages (using AWS API Gateway). The three
services are the following:

  * •

Catalog service: It manages the operations related with the products.

  * •

Customer service: It manages everything related with customers.

  * •

Order service: It involves all in relation to orders.

The scheme of the system will be better explained in the Figure 1. The
performance of this initial solution can be seen in Table I with an average
latency of 512 ms when lambdas are warmed and over 16307 ms on cold starts. It
is important to highlight that the test lambda functions are mainly based on
interaction with DynamoDB and are not computationally intensive.

Initial Solution (warmed functions)
---
functions / workload | 1000 users | 5000 users | 10000 users
customer lambda | 482 ms | 383 ms | 390 ms
catalog lambda | 564 ms | 487 ms | 495 ms
order lambda | 657 ms | 589 ms | 581 ms
TABLE I: Performance measurement of initial solution

##  III Performance tests

In line with the main objective of the study, load tests will be conducted for
each of the implemented improvements. The objectives of these tests is
quantify the improvement offered by each of the implemented techniques or
approaches and compare the latencies with the initial unoptimized system.
Due to that, the tests have a duration of 30 minutes in which the tested
system is exposed to different constant workloads. Each test is repeated a
total of 3 times, with the result being the average value of all repetitions.
To evaluate cold start performance, functions were invoked manually from the
AWS console, and the experienced times were collected. A test is compound of 3
requests, one for each function of the application. The results will be
measured and displayed in milliseconds (ms).

##  IV Best practices and techniques

Below are the proposed techniques and approaches to reduce the cold start of
Java functions and improve their overall performance. Along with their
measurements in terms of performance tests. The proposed practices will be
compared in the VI.Discussion section.

###  IV-A Better suited configuration

Increasing the allocated memory for a function is probably the most well-known
way to reduce cold starts and also the simplest. The allocation of memory to a
function is related to the amount of CPU assigned to that function. A greater
amount of CPU generally makes the initialization of the function faster.
However, it’s not always the case that more memory equals better performance.
It’s also necessary to consider that higher memory allocation will result in
higher costs, but if the computing time is reduced due to increased
performance, the overall cost will decrease as well. To properly balance the
memory configuration of a Lambda function, the Lambda Power Tuner solution
[19] can be of great help to profile the function, displaying graphs such as
the one shown in Figure 2 for Catalog lambda of our solution.

Figure 2: Power tuner with Catalog lambda results

In our case, chose the right memory value for the functions configuration
suppose an improvement of about 70% in terms of cold start going from an
average of 16307 ms to 4867 ms. In addition to a noteworthy 46% when the
functions are warmed passing from an average of 512 ms to 216 ms as can be
seen in Table II. For the rest of the versions that do not have an appropriate
memory configuration, the default value of 512 MB is used.

Configuration optimized solution (warmed functions)
---
functions / workload | 1000 users | 5000 users | 10000 users
customer lambda | 305 ms | 244 ms | 257 ms
catalog lambda | 321 ms | 249 ms | 264 ms
order lambda | 312 ms | 258 ms | 277 ms
TABLE II: Performance measurement of configuration optimized solution

###  IV-B Snapstart

Snapstart [16] is a feature offered by AWS to mitigate the cold start of Java
functions. It involves initializing the function at the moment a version of it
is created. Lambda takes a snapshot of the initialized lambda’s runtime
environment, encrypts and caches it. When initializing new functions, it does
so from the snapshot rather than starting from scratch, thus improving
performance.
A notable drawback of this feature is that it does not support many other
features offered by AWS Lambda, such as custom runtimes, the arm64
architecture, Amazon EFS, and other managed runtimes that are not Java 11 or
Java 17.
In this case, enable Snapstart on our functions allow an enhancement of 16% at
the cold start (going from an average of 16307 ms to 13736 ms) and an
improvement of 21% attending requests when the functions are warmed making a
difference of over 100 ms in each request answered (Table III).

Snapstart (warmed functions)
---
functions / workload | 1000 users | 5000 users | 10000 users
customer lambda | 427 ms | 306 ms | 364 ms
catalog lambda | 488 ms | 271 ms | 410 ms
order lambda | 560 ms | 282 ms | 525 ms
TABLE III: Performance measurement of Snapstart solution

###  IV-C Arm64 architecture

The instruction set architecture within AWS Lambda, either x86_64 or arm64,
determines the processor that the function will use. AWS’s official
documentation suggests that the arm64 architecture [18] with its Graviton
processor can offer significantly better performance than the x86_64
architecture.
Regarding our solution, the impact of using arm64 on the latencies was
positive, getting over a 14% of improvement on cold starts and a 13% when the
lambdas are warmed. Reducing the average latency by over 2000 ms and
approximately 70 ms, respectively. The warmed lambda functions data can be
seen in Table IV.

arm64 (warmed functions)
---
functions / workload | 1000 users | 5000 users | 10000 users
customer lambda | 374 ms | 376 ms | 355 ms
catalog lambda | 455 ms | 440 ms | 418 ms
order lambda | 542 ms | 543 ms | 506 ms
TABLE IV: Performance measurement of arm64 solution

###  IV-D AWS SDK v2 for Java

The AWS SDK for Java 2.x [15] is the second version of the SDK provided by AWS
for Java. This version is designed with a modular approach, allowing you to
include only the service-specific JARs you need, reducing the size of your
artifacts. That’s a key point of reducing the cold start, loading less number
of classes when initialize the function. Furthermore, it boasts several other
performance improvements such as asynchronous nature, connection management,
concurrency support, resource management, and so forth. Due to that, in terms
of performance its use is highly recommended.
For us, migrate the lambda functions from sdk v1 to sdk v2 suppose an
enhancement of about 40% in both cold start and warmed functions. On first
going from 16307 ms to 9926 ms and on second one passing from 512 ms to 298 ms
of latency as is visible in Table V.

sdk v2 (warmed functions)
---
functions / workload | 1000 users | 5000 users | 10000 users
customer lambda | 307 ms | 277 ms | 267 ms
catalog lambda | 300 ms | 290 ms | 279 ms
order lambda | 340 ms | 317 ms | 305 ms
TABLE V: Performance measurement of sdkv2 solution

###  IV-E GraalVM

GraalVM [20] is a platform that enables ahead-of-time compilation of Java
code, leveraging the Native Image tool to produce a standalone binary. As a
result, this binary offers significantly better performance compared to
traditional Java code that runs using the JVM and is compiled at runtime. The
binary is produced for the OS of the environment the code is built. It’s
important to note that AWS Lambda is only compatible with GraalVM through the
use of custom runtimes. Other disadvantage is that the ahead-of-time
compilation sometimes would be tricky to make it work properly adding even the
necessity of dealing with low level issues caused by AOT compilation.
For us, create a native image of our code and uses it generates a great
improvement of 83% on cold start, going form 16307 ms to 2800 ms. When the
lambdas are warmed, it goes from and average of 512 ms to 231 ms which means a
55% of enhancement, visible in Table VI.

GraalVM (warmed functions)
---
functions / workload | 1000 users | 5000 users | 10000 users
customer lambda | 269 ms | 193 ms | 231 ms
catalog lambda | 272 ms | 194 ms | 229 ms
order lambda | 265 ms | 196 ms | 230 ms
TABLE VI: Performance measurement of GraalVM solution

###  IV-F Environment variable: JAVA_TOOL_OPTIONS

AWS Lambda provides some customization options [17] for the managed Java
runtime, which can enhance overall function performance. This is achieved
through the use of the JAVA_TOOL_OPTIONS environment variable, which allows
access to JVM features such as tiered compilation or garbage collector
behavior among others. In this case we use only tiered compilation, fixing the
level to 1, enabling the C1 compiler which means optimized code for fast
start-up time. The exact value for the mentioned key is
”-XX:+TieredCompilation -XX:TieredStopAtLevel=1”. Java 17 managed runtime is
default set to level 1 for AWS Lambda, while in Java 11, you need to set it
manually. Level two can be used for an overall better performance, not only
focus on the start time. In our test system, enabling that option translates
to a 40% improvement in cold start times and a 31% improvement in response
times for the system with warm Lambdas (Table VII).

JAVA_TOOL_OPTIONS (warmed functions)
---
functions / workload | 1000 users | 5000 users | 10000 users
customer lambda | 359 ms | 311 ms | 327 ms
catalog lambda | 368 ms | 336 ms | 350 ms
order lambda | 402 ms | 369 ms | 365 ms
TABLE VII: Performance measurement of java_options solution

###  IV-G Others

In addition to the previous techniques and approaches we propose and have
tested, there are many other best practices [14, 21, 6]. Whether they are more
dependent on the use case or more general, below, we mention some of them:

  * •

Reduce the size of the bundle to enhance the initialization phase. As less
classes to load as fast the function will initialize.

  * •

Initialize the connections of third party out of the handler and same the
static resources on /tmp to reuse it from one execution to another.

  * •

A framework, if needed, should be as lighter as possible. Due to that we
suggest replacing spring with other smaller frameworks like Vert.x

  * •

Try to avoid reflection, that feature makes some JVM optimizations impossible
to be done as it works with dynamic types.

  * •

Initialize all necessary dependencies and classes during initialization time.

These approaches, combined with some of the previous ones we have tested, for
sure will greatly reduce the impact of cold starts and increase the
performance of Java functions by a really significant percentage.

##  V Combinations

After defining the proposed techniques and quantifying the effects of their
application, it becomes essential to attempt to combine them with the
intention of detecting synergies, incompatibilities, and understanding the
improvements that can be achieved by applying several of these. With this
intention, and due to the unavailability of Snapstart in combination with the
Arm64 architecture, the following combinations are proposed:

  * •

Best memory configuration + SDK v2 + JAVA_TOOL_OPTIONS environment variable +
Snapstart

  * •

Best memory configuration + SDK v2 + JAVA_TOOL_OPTIONS environment variable +
Arm64 architecture

Two new prototypes have been created, optimized with the previous optimization
combinations, and their performance has been measured. The results are
compared with each other and independently with the GraalVM prototype.

Figure 3: Percentage of improvement by combination %

In Figure 3, you can observe the percentage of improvement compared to the
performance offered by the initial system, both for the cold start of
functions and for the overall performance.

##  VI Discussion

In Figure 4 and Figure 5, you can see a summary of the cold start improvement
and warm Lambdas, respectively. Next, we will analyze each approach and the
possibilities they offer.
Starting with an appropriate configuration for each function, in our opinion,
is the minimum optimization that any function should have, both to improve
performance and to balance cost. In addition to offering a significant
improvement, it is perfectly compatible with all the other techniques
mentioned in this article.

Figure 4: Cold start enhancement by approach %

Continuing with the use of Snapstart, this technique focuses on mitigating
cold starts. In our experience, this technique could not produce as
significant an improvement as to justify the incompatibilities it generates,
including the lack of support for the arm64 architecture or custom runtimes,
along with other important features such as the capability to use EFS or to
attach the lambda to a VPC.
The improvement offered by using the Arm64 architecture seems to depend on the
specific use case in which it is employed. For example, it can be particularly
beneficial for compute-intensive applications such as high-performance
computing. This could explain why the improvement in our case is not too
noticeable. Nevertheless, its use is highly recommended as it can be combined
with other approaches discussed, and its adoption is on the rise. The major
drawback is the incompatibility with some third-party technologies or
dependencies related to this architecture.

Figure 5: Warmed function enhancement by approach %

The use of AWS’s SDK v2 for Java applications is also practically mandatory,
as long as it supports all the services and the majority of the dependencies
integrated into the function. Thanks to its improvements and modular approach,
this second version of the SDK significantly increases performance. Moreover,
it is perfectly mixable with the rest of the techniques without adding notable
limitations. Combined with version 1, if any unsupported features are needed,
they can be utilized as required.
As expected, leveraging the benefits of ahead-of-time compilation generates a
significant performance boost both in cold starts and when functions are
already warm. However, the main drawback is the complexity that this approach
adds. It’s common to encounter runtime issues that may require the developer
to deal with low-level aspects, and any modifications can become challenging.
In principle, it is perfectly combinable with the rest of the techniques that
support custom runtimes.
The possibilities offered by the JAVA_TOOL_OPTIONS environment variable are
quite extensive, allowing for significant Java VM configuration. In our case,
we only tested tiered compilation, but there are other possibilities such as
configuring garbage collector behavior [7]. In any case, it’s worth trying to
use this variable and integrating it into Java functions whenever the JVM is
used because it doesn’t impose any limitations, and the improvement is
significant.
It’s also important to highlight that investing efforts in refactoring the
function’s code and applying the details mentioned in subsection IV-G can
achieve an even greater performance improvement than the techniques discussed.
This work is entirely dependent on the use case and requires trial and error
as well as skill.
Regarding the simultaneous application of techniques, for our combinations, we
recommend the use of either of the two, so you should choose the one that
provides the most benefit, either in terms of performance or greater
adaptability to the specific use case.

##  VII Related work

As far as we know, and up to this date, there aren’t too many articles seeking
to enhance the Java approach with AWS Lambda in the same manner as we do.
Nonetheless, there are some similar studies. In [21] Buliani addresses part of
the performance issue and proposes and applies a series of techniques
incrementally, analyzing the overall improvement step by step. [12]
investigates the performance impact of cold start and containerization for
Java-based FaaS functions and compares it to an alternative using Python. [2]
takes a similar approach, suggesting a variety of ways to enhance the
performance of AWS Lambda but without integrating it with Java. In [5] Dowd
investigates the overhead and cold start for functions written in Java,
attempting to determine the extent to which the findings depend on the
programming language.
There are also other studies that have a certain relationship with ours,
either because they explore some of the techniques we propose or because they
are related with the performance of AWS Lambda. [3, 4] analyze the performance
of the various managed runtimes provided by AWS for executing code in
different languages on Lambda functions without optimizing it. In [1]
Hussachai et al demonstrates how optimizing the Java artifact results in
superior performance during the execution of the Lambda function. In [7]
evaluate the use of a garbage collector control strategy concerning
performance improvement, that is an interesting strategy that we didn’t take
in considerationin our study. Related with GraalVM, [10, 8] incorporate this
tool with the aim of leveraging the advantage it offers in terms of
performance.

##  VIII Conclusion

Performance is an aspect that should always be taken into account when
creating Lambda functions, especially if you plan to combine them with Java.
Once the performance improvements offered by these approaches for using Java
with AWS Lambda have been quantified. It is evident that Java can be used
competitively with AWS Lambda without the need to resort to runtimes like
Node.js or Python (which usually report better performance).
Having tested various combinations of proposed strategies and demonstrated
that the performance is similar to what can be achieved with the use of
GraalVM, we recommend avoiding its usage. We consider it essential to profile
the Java functions to select the optimal memory configuration and using AWS
SDK v2. From there, optionally, the next step is to the rest of the techniques
to achieve one of the proposed combinations. Additionally, following the
general best practices such as initializing connections outside the handler or
reduce the bundle size, among others, also contribute to optimizing Java
functions.
The optimization strategies we propose are not the only ones available, there
are many other techniques that can similarly yield excellent results.
It is important to highlight that the process of optimizing a function is
linear in terms of difficulty. In other words, it’s relatively easy to obtain
a significant improvement, but if you want to squeeze out the maximum
performance, the process requires trial and error and greater effort. It’s
also important to highlight that the percentage of improvement offered by each
approach depends on the specific use case in which they are applied. This
percentage can vary, either increasing or decreasing.

## References

  * [1] Hussachai Puripunpinyo and M.H. Samadzadeh  “Effect of optimizing Java deployment artifacts on AWS Lambda”  In _2017 IEEE Conference on Computer Communications Workshops (INFOCOM WKSHPS)_ , 2017, pp. 438–443  DOI: 10.1109/INFCOMW.2017.8116416
  * [2] Daniel Bardsley, Larry Ryan and John Howard  “Serverless Performance and Optimization Strategies”  In _2018 IEEE International Conference on Smart Cloud (SmartCloud)_ , 2018, pp. 19–26  DOI: 10.1109/SmartCloud.2018.00012
  * [3] David Jackson and Gary Clynch  “An Investigation of the Impact of Language Runtime on the Performance and Cost of Serverless Functions”  In _2018 IEEE/ACM International Conference on Utility and Cloud Computing Companion (UCC Companion)_ , 2018, pp. 154–160  DOI: 10.1109/UCC-Companion.2018.00050
  * [4] Omid Hosseini  “AWS Lambda Language Performance”, 2019  URL: https://gupea.ub.gu.se/handle/2077/62454
  * [5] Abigail Lauren Dowd  “Performance of Java in Function-as-a-Service computing”  In _University of Texas Libraries_ , 2020  URL: https://repositories.lib.utexas.edu/handle/2152/84357
  * [6] Sean O’Toole  “AWS Lambda Java tutorial: Reducing cold starts”, 2020  URL: https://www.capitalone.com/tech/cloud/aws-lambda-java-tutorial-reduce-cold-starts/?v=1694091771423
  * [7] David Quaresma, Daniel Fireman and Thiago Emmanuel Pereira  “Controlling Garbage Collection and Request Admission to Improve Performance of FaaS Applications”  In _2020 IEEE 32nd International Symposium on Computer Architecture and High Performance Computing (SBAC-PAD)_ , 2020, pp. 175–182  DOI: 10.1109/SBAC-PAD49847.2020.00033
  * [8] M. Šipek, D. Muharemagić, B. Mihaljević and A. Radovan  “Enhancing Performance of Cloud-based Software Applications with GraalVM and Quarkus”  In _2020 43rd International Convention on Information, Communication and Electronic Technology (MIPRO)_ , 2020, pp. 1746–1751  DOI: 10.23919/MIPRO48935.2020.9245290
  * [9] Qusay I. Hassan B.  “Survey on serverless computing”  In _Journal of Cloud Computing_ , 2021
  * [10] Carlos David Ron Mosquera  “Mejora del rendimiento de aplicaciones serverless basadas en la nube con GraalVM y Quarkus”, 2022
  * [11] StackOverflow  “Developer survey 2022”, 2022  URL: https://survey.stackoverflow.co/2022/
  * [12] Qinzhe Wu and Lizy K. John  “Performance of Java in Function-as-a-Service Computing”  In _2022 IEEE/ACM 15th International Conference on Utility and Cloud Computing (UCC)_ , 2022, pp. 261–266  DOI: 10.1109/UCC56403.2022.00047
  * [13] Datdog  “The state of serverless”  In _datadog_ , 2023  URL: https://www.datadoghq.com/state-of-serverless/
  * [14] AWS  “Best practices to work with AWS Lambda functions”  URL: https://docs.aws.amazon.com/es_es/lambda/latest/dg/best-practices.html
  * [15] AWS  “Developer Guide - AWS SDK for Java 2.x”  URL: https://docs.aws.amazon.com/sdk-for-java/latest/developer-guide/home.html
  * [16] AWS  “Improving startup performance with Lambda SnapStart”  URL: https://docs.aws.amazon.com/lambda/latest/dg/snapstart.html
  * [17] AWS  “Java Lambda functions customization settings”  URL: https://docs.aws.amazon.com/lambda/latest/dg/java-customization.html
  * [18] AWS  “Lambda instruction set architectures (ARM/x86)”  URL: https://docs.aws.amazon.com/es_es/lambda/latest/dg/foundation-arch.html
  * [19] AWS  “Profiling functions with AWS Lambda Power Tuning”  URL: https://docs.aws.amazon.com/lambda/latest/operatorguide/profile-functions.html
  * [20] GraalVM org  “GraalVM overview”  URL: https://www.graalvm.org/latest/docs/introduction/
  * [21] AWS Stefano Buliani  “Best practices for AWS Lambda and Java”  URL: https://d1.awsstatic.com/events/reinvent/2019/REPEAT_1_Best_practices_for_AWS_Lambda_and_Java_SVS403-R1.pdf