Original: D YNAMIC VAE: D ECOUPLING R ECONSTRUCTION E RROR AND D ISENTANGLED R EPRESENTATION L EARNING
Revised: D YNAMIC VAE: D ECOUPLING R ECONSTRUCTION E RROR AND D ISENTANGLED R EPRESENTATION L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: This paper challenges the common assumption that the weight β , in β -VAE, should be larger than 1 in order to effectively disentangle latent factors.
Revised: This paper challenges the common assumption that the weight β , in β -VAE, should be larger than 1 in order to effectively disentangle latent factors.

Original: We demonstrate that β -VAE, with β < 1 , can not only attain good disentanglement but also signiﬁcantly improve reconstruction accuracy via dynamic control.
Revised: We demonstrate that β -VAE, with β < 1 , can not only attain good disentanglement but also signiﬁcantly improve reconstruction accuracy via dynamic control.

Original: The paper removes the inherent trade-off between reconstruction accuracy and disentanglement for β -VAE.
Revised: The paper removes the inherent trade-off between reconstruction accuracy and disentanglement for β -VAE.

Original: Existing methods, such as β -VAE and FactorVAE, assign a large weight to the KL-divergence term in the objective function, leading to high reconstruction errors for the sake of better disentanglement.
Revised: Existing methods, such as β -VAE and FactorVAE, assign a large weight to the KL-divergence term in the objective function, leading to high reconstruction errors for the sake of better disentanglement.

Original: To mitigate this problem, a ControlVAE has recently been developed that dynamically tunes the KL-divergence weight in an attempt to control the trade-off to more a favorable point.
Revised: To mitigate this problem, a ControlVAE has recently been developed that dynamically tunes the KL-divergence weight in an attempt to control the trade-off to more a favorable point.

Original: However,
Revised: However,

Original: ControlVAE fails to eliminate the conﬂict between the need for a large β (for disentanglement) and the need for a small β (for smaller reconstruction error).
Revised: ControlVAE fails to eliminate the conﬂict between the need for a large β (for disentanglement) and the need for a small β (for smaller reconstruction error).

Original: Instead, we propose DynamicVAE that maintains a different β at different stages of training, thereby decoupling disentanglement and reconstruction accuracy .
Revised: Instead, we propose DynamicVAE that maintains a different β at different stages of training, thereby decoupling disentanglement and reconstruction accuracy .

Original: In order to evolve the weight, β , along a trajectory that enables such decoupling, DynamicVAE leverages a modiﬁed incremental PI (proportional-integral) controller, a variant of proportional-integral-derivative controller (PID) algorithm, and employs a moving average as well as a hybrid annealing method to evolve the value of
Revised: In order to evolve the weight, β , along a trajectory that enables such decoupling, DynamicVAE leverages a modiﬁed incremental PI (proportional-integral) controller, a variant of proportional-integral-derivative controller (PID) algorithm, and employs a moving average as well as a hybrid annealing method to evolve the value of

Original: KL-divergence smoothly in a tightly controlled fashion.
Revised: KL-divergence smoothly in a tightly controlled fashion.

Original: We theoretically prove the stability of the proposed approach.
Revised: We theoretically prove the stability of the proposed approach.

Original: Evaluation results on three benchmark datasets demonstrate that DynamicVAE signiﬁcantly improves the reconstruction accuracy while achieving disentanglement comparable to the best of existing methods.
Revised: Evaluation results on three benchmark datasets demonstrate that DynamicVAE signiﬁcantly improves the reconstruction accuracy while achieving disentanglement comparable to the best of existing methods.

Original: The results verify that our method can separate disentangled representation learning and reconstruction, removing the inherent tension between the two.
Revised: The results verify that our method can separate disentangled representation learning and reconstruction, removing the inherent tension between the two.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: The goal of disentangled representation learning is to encode input data into a low-dimensional space that preserves information about the salient factors of variation, so that each dimension of the representation corresponds to a distinct factor in the data (Bengio et al., 2013; Locatello et al., 2020; van Steenkiste et al., 2019).
Revised: The goal of disentangled representation learning is to encode input data into a low-dimensional space that preserves information about the salient factors of variation, so that each dimension of the representation corresponds to a distinct factor in the data (Bengio et al., 2013; Locatello et al., 2020; van Steenkiste et al., 2019).

Original: Learning disentangled representations beneﬁts a variety of downstream tasks (Higgins et al., 2018;
Revised: Learning disentangled representations beneﬁts a variety of downstream tasks (Higgins et al., 2018;

Original: Lake et al., 2017; Locatello et al., 2019c;a; Denton et al., 2017; Mathieu et al., 2019), including abstract visual reasoning (van Steenkiste et al., 2019),
Revised: Lake et al., 2017; Locatello et al., 2019c;a; Denton et al., 2017; Mathieu et al., 2019), including abstract visual reasoning (van Steenkiste et al., 2019),

Original: zero-shot transfer learning (Burgess et al., 2018; Lake et al., 2017; Higgins et al., 2017a) and image generation (Nie et al., 2020), just to name a few.
Revised: zero-shot transfer learning (Burgess et al., 2018; Lake et al., 2017; Higgins et al., 2017a) and image generation (Nie et al., 2020), just to name a few.

Original: Due to its central importance in various downstream applications, there is abundant literature on learning disentangled representations.
Revised: Due to its central importance in various downstream applications, there is abundant literature on learning disentangled representations.

Original: Roughly speaking, there are two lines of methods towards this goal.
Revised: Roughly speaking, there are two lines of methods towards this goal.

Original: The ﬁrst category includes supervised methods (Chen &
Revised: The ﬁrst category includes supervised methods (Chen &

Original: Batmanghelich, 2019; Locatello et al., 2019c; Shu et al., 2019; Bouchacourt et al., 2018; Nie et al., 2020; Yang et al., 2015), where external supervision (e.g., data generative factors) is available during training to guide the learning of disentangled representations.
Revised: Batmanghelich, 2019; Locatello et al., 2019c; Shu et al., 2019; Bouchacourt et al., 2018; Nie et al., 2020; Yang et al., 2015), where external supervision (e.g., data generative factors) is available during training to guide the learning of disentangled representations.

Original: The second line of works focus on unsupervised methods (Chen et al., 2016; 2018; Burgess et al., 2018; Kim & Mnih, 2018; Denton et al., 2017; Kumar et al., 2018; Fraccaro et al., 2017), which substantially relieve the needs to have external supervisions.
Revised: The second line of works focus on unsupervised methods (Chen et al., 2016; 2018; Burgess et al., 2018; Kim & Mnih, 2018; Denton et al., 2017; Kumar et al., 2018; Fraccaro et al., 2017), which substantially relieve the needs to have external supervisions.

Original: For this reason, in this paper, we mainly focus on unsupervised disentangled representation learning.
Revised: For this reason, in this paper, we mainly focus on unsupervised disentangled representation learning.

Original: One major challenge of unsupervised disentanglement learning is that there exists a trade-off between reconstruction quality of the input signal and the degree of disentanglement in the latent representations.
Revised: One major challenge of unsupervised disentanglement learning is that there exists a trade-off between reconstruction quality of the input signal and the degree of disentanglement in the latent representations.

Original: Let us take β -VAE and its variants (Burgess et al., 2018; Chen et al., 2018; Higgins et al., 2017a) as an example.
Revised: Let us take β -VAE and its variants (Burgess et al., 2018; Chen et al., 2018; Higgins et al., 2017a) as an example.

Original: These methods assign a large and ﬁxed weight β in the objective function to improve the disentanglement at the cost of reconstruction quality, which is highly correlated with accuracy in downstream tasks (van Steenkiste et al., 2019; Locatello et al., 2020).
Revised: These methods assign a large and ﬁxed weight β in the objective function to improve the disentanglement at the cost of reconstruction quality, which is highly correlated with accuracy in downstream tasks (van Steenkiste et al., 2019; Locatello et al., 2020).

Original: In order to improve the reconstruction quality, researchers have proposed a dynamic learning approach, ControlVAE (Shao et al., 2020), to dynamically adjust the weight on the KL term in the VAE objective to better balance the quality of disentangled representation learning and reconstruction error.
Revised: In order to improve the reconstruction quality, researchers have proposed a dynamic learning approach, ControlVAE (Shao et al., 2020), to dynamically adjust the weight on the KL term in the VAE objective to better balance the quality of disentangled representation learning and reconstruction error.

Original: However, while ControlVAE allows better control of the trade-off between disentangled representation learning and reconstruction error, it does not eliminate it .
Revised: However, while ControlVAE allows better control of the trade-off between disentangled representation learning and reconstruction error, it does not eliminate it .

Original: One is still achieved at the expense of the other.
Revised: One is still achieved at the expense of the other.

Original: The contribution of this paper, compared to the above state of the art, lies in demonstrating that with the proper design, the trade-off between disentangled representation learning and reconstruction error is completely eliminated.
Revised: The contribution of this paper, compared to the above state of the art, lies in demonstrating that with the proper design, the trade-off between disentangled representation learning and reconstruction error is completely eliminated.

Original: Both objectives can be attained at the same time in a decoupled fashion, without affecting each other.
Revised: Both objectives can be attained at the same time in a decoupled fashion, without affecting each other.

Original: More speciﬁcally, we observe that if β was kept high in the beginning of training then lowered later in the process, the two objectives are decoupled allowing each to be independently optimized.
Revised: More speciﬁcally, we observe that if β was kept high in the beginning of training then lowered later in the process, the two objectives are decoupled allowing each to be independently optimized.

Original: To the authors’ knowledge, this work is the ﬁrst to attain such decoupled optimization of both quality of disentanglement and reconstruction error.
Revised: To the authors’ knowledge, this work is the ﬁrst to attain such decoupled optimization of both quality of disentanglement and reconstruction error.

Original: Our Contributions: In this paper, we propose a novel unsupervised disentangled representation learning method, dubbed as DynamicVAE, that turns the weight of β -VAE ( β > 1 ) (Burgess et al., 2018; Higgins et al., 2017a) into a small value ( β ≤ 1 ) to achieve not only good disentanglement but also a high reconstruction accuracy via dynamic control.
Revised: Our Contributions: In this paper, we propose a novel unsupervised disentangled representation learning method, dubbed as DynamicVAE, that turns the weight of β -VAE ( β > 1 ) (Burgess et al., 2018; Higgins et al., 2017a) into a small value ( β ≤ 1 ) to achieve not only good disentanglement but also a high reconstruction accuracy via dynamic control.

Original: We summarize the main contributions of this paper as follows
Revised: We summarize the main contributions of this paper as follows

Original: • We propose a new model, DynamicVAE, that leverages an incremental PI controller and moving average to evolve the desired KL-divergence along a trajectory that enables decoupling of two objectives : high-quality disentanglement and low reconstruction error.
Revised: • We propose a new model, DynamicVAE, that leverages an incremental PI controller and moving average to evolve the desired KL-divergence along a trajectory that enables decoupling of two objectives : high-quality disentanglement and low reconstruction error.

Original: • We provide the theoretical conditions on parameters of the PI controller to guarantee stability of DynamicVAE.
Revised: • We provide the theoretical conditions on parameters of the PI controller to guarantee stability of DynamicVAE.

Original: • We experimentally demonstrate that our approach turns the weight of β -VAE ( β > 1 ) to β ≤ 1 , achieving higher reconstruction quality yet comparable disentanglement compared to prior approaches (e
Revised: • We experimentally demonstrate that our approach turns the weight of β -VAE ( β > 1 ) to β ≤ 1 , achieving higher reconstruction quality yet comparable disentanglement compared to prior approaches (e

Original: g., FactorVAE).
Revised: g., FactorVAE).

Original: Thus, our results verify that the proposed method indeed decouples disentanglement and reconstruction accuracy without hurting each other’s performance.
Revised: Thus, our results verify that the proposed method indeed decouples disentanglement and reconstruction accuracy without hurting each other’s performance.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: β -VAE and its Variants: β -VAE (Higgins et al., 2017b;
Revised: β -VAE and its Variants: β -VAE (Higgins et al., 2017b;

Original: Chen et al., 2018) is a popular unsupervised method for learning disentangled representations of the data generative factors (Bengio et al., 2013).
Revised: Chen et al., 2018) is a popular unsupervised method for learning disentangled representations of the data generative factors (Bengio et al., 2013).

Original: Compared to the original VAE, β -VAE incorporates an extra hyperparameter β ( β > 1) as the weight of the KL term in the VAE objective:
Revised: Compared to the original VAE, β -VAE incorporates an extra hyperparameter β ( β > 1) as the weight of the KL term in the VAE objective:

Original: [Equation]
Revised: [Equation]

Original: In order to discover more disentangled factors, in other variants, practitioners further add a constraint on the total information capacity, C , to control the capacity of the latent channels (Burgess et al., 2018) to transmit information.
Revised: In order to discover more disentangled factors, in other variants, practitioners further add a constraint on the total information capacity, C , to control the capacity of the latent channels (Burgess et al., 2018) to transmit information.

Original: The constraint can be formulated as an optimization method:
Revised: The constraint can be formulated as an optimization method:

Original: [Equation]
Revised: [Equation]

Original: where β is a large and ﬁxed hyperparameter.
Revised: where β is a large and ﬁxed hyperparameter.

Original: As a result, when the weight β is large (e.g. 100), the algorithm tends to optimize the second term in (2), leading to much higher reconstruction error.
Revised: As a result, when the weight β is large (e.g. 100), the algorithm tends to optimize the second term in (2), leading to much higher reconstruction error.

Original: PID Control Algorithm:
Revised: PID Control Algorithm:

Original: The PID is a simple yet effective control algorithm that can stabilize system output to a desired value via feedback control (Stooke et al., 2020; ˚Astr¨om et al., 2006).
Revised: The PID is a simple yet effective control algorithm that can stabilize system output to a desired value via feedback control (Stooke et al., 2020; ˚Astr¨om et al., 2006).

Original: The
Revised: The

Original: PID algorithm calculates an error, e ( t ) , between a set point (in this case, the desired KL-divergence) and the current value of the controlled variable (in this case, the actual KL-divergence), then applies a correction in a direction that reduces that error.
Revised: PID algorithm calculates an error, e ( t ) , between a set point (in this case, the desired KL-divergence) and the current value of the controlled variable (in this case, the actual KL-divergence), then applies a correction in a direction that reduces that error.

Original: The correction is the weighted sum of three terms, one proportional to the error (called P), one that is the integral of error (called I), and one that is the derivative of error (called D); thus, the term PID.
Revised: The correction is the weighted sum of three terms, one proportional to the error (called P), one that is the integral of error (called I), and one that is the derivative of error (called D); thus, the term PID.

Original: The derivative term is not recommended for noisy systems, such as ours, reducing the algorithm to PI control.
Revised: The derivative term is not recommended for noisy systems, such as ours, reducing the algorithm to PI control.

Original: The canonical form of a PI controller
Revised: The canonical form of a PI controller

Original: (applied to control β ( t ) ) is the following:
Revised: (applied to control β ( t ) ) is the following:

Original: [Equation]
Revised: [Equation]

Original: where β ( t ) is the output of a controller, which (in our case) is the used β during training at time t ; e ( t ) is the error between the output value and the desired value at time t ; K p , K i denote the coefﬁcients for the P term and I term, respectively.
Revised: where β ( t ) is the output of a controller, which (in our case) is the used β during training at time t ; e ( t ) is the error between the output value and the desired value at time t ; K p , K i denote the coefﬁcients for the P term and I term, respectively.

Original: Eq.
Revised: Eq.

Original: (3) may be rewritten in incremental form, as follows:
Revised: (3) may be rewritten in incremental form, as follows:

Original: [Equation]
Revised: [Equation]

Original: where β (0) can be set as needed (as we show later), and:
Revised: where β (0) can be set as needed (as we show later), and:

Original: [Equation]
Revised: [Equation]

Original: This paper adopts a nonlinear incremental form of the PI controller, described later in Section 3.
Revised: This paper adopts a nonlinear incremental form of the PI controller, described later in Section 3.

Original:  T HE D YNAMIC VAE A LGORITHM
Revised:  T HE D YNAMIC VAE A LGORITHM

Original: The goal of disentangled representation learning (Burgess et al., 2018) is to maximize the log likelihood and simultaneously stabilize the KL-divergence to a target value C .
Revised: The goal of disentangled representation learning (Burgess et al., 2018) is to maximize the log likelihood and simultaneously stabilize the KL-divergence to a target value C .

Original: It can be formulated as the following constrained optimization problem:
Revised: It can be formulated as the following constrained optimization problem:

Original: [Equation]
Revised: [Equation]

Original: In order to achieve a good trade-off between disentanglement and reconstruction accuracy, we attempt to design a controller to dynamically adjust β ( t ) in the following VAE objective to stabilize the
Revised: In order to achieve a good trade-off between disentanglement and reconstruction accuracy, we attempt to design a controller to dynamically adjust β ( t ) in the following VAE objective to stabilize the

Original: KL-divergence to the desired value C :
Revised: KL-divergence to the desired value C :

Original: [Equation]
Revised: [Equation]

Original: The contribution of DynamicVAE is to evolve β ( t ) along a good trajectory to achieve decoupling between disentanglement
Revised: The contribution of DynamicVAE is to evolve β ( t ) along a good trajectory to achieve decoupling between disentanglement

Original: and reconstruction error
Revised: and reconstruction error

Original: To reach this goal, we need to address the following two challenges:
Revised: To reach this goal, we need to address the following two challenges:

Original: β ( t ) should dynamically change from a large value to small one.
Revised: β ( t ) should dynamically change from a large value to small one.

Original: Speciﬁcally, at the beginning of training, β ( t ) should be large enough to disentangle latent factors.
Revised: Speciﬁcally, at the beginning of training, β ( t ) should be large enough to disentangle latent factors.

Original: After that, β ( t ) is required to gradually drop to a small value to optimize the reconstruction.
Revised: After that, β ( t ) is required to gradually drop to a small value to optimize the reconstruction.

Original: β ( t ) should not change too fast or oscillates too frequently.
Revised: β ( t ) should not change too fast or oscillates too frequently.

Original: When β ( t ) drops too fast or oscillates, it may cause KL-divergence to grow with a large value.
Revised: When β ( t ) drops too fast or oscillates, it may cause KL-divergence to grow with a large value.

Original: Consequently, some latent factors may emerge earlier so that they can potentially be entangled with each other.
Revised: Consequently, some latent factors may emerge earlier so that they can potentially be entangled with each other.

Original: In this paper, we propose methods to deal with these two challenges, summarized below.
Revised: In this paper, we propose methods to deal with these two challenges, summarized below.

Original: A non-linear incremental PI controller: Fig.
Revised: A non-linear incremental PI controller: Fig.

Original: 1 (a) shows the designed non-linear PI controller that dynamically adjusts the weight β ( t ) in the KL term of the β -VAE, based on the actual KL-divergence, y KL ( t ) .
Revised: 1 (a) shows the designed non-linear PI controller that dynamically adjusts the weight β ( t ) in the KL term of the β -VAE, based on the actual KL-divergence, y KL ( t ) .

Original: Speciﬁcally, it ﬁrst samples the output KL-divergence, y KL ( t ) , at training step t .
Revised: Speciﬁcally, it ﬁrst samples the output KL-divergence, y KL ( t ) , at training step t .

Original: Then we use the difference e ( t ) between the sampled KL-divergence at time t with the desired value, C , as the feedback to PI controller to tune β ( t ) .
Revised: Then we use the difference e ( t ) between the sampled KL-divergence at time t with the desired value, C , as the feedback to PI controller to tune β ( t ) .

Original: The corresponding PI algorithm is denoted by
Revised: The corresponding PI algorithm is denoted by

Original: [Equation]
Revised: [Equation]

Original: where σ ( . ) is a sigmoid function; K p , K i are positive hyper-parameters for P and I term respectively.
Revised: where σ ( . ) is a sigmoid function; K p , K i are positive hyper-parameters for P and I term respectively.

Original: As mentioned earlier, we need a large β ( t ) in the beginning to control the KL-divergence from a small value to a large target value so that the information can be transmitted through the latent channels per data sample.
Revised: As mentioned earlier, we need a large β ( t ) in the beginning to control the KL-divergence from a small value to a large target value so that the information can be transmitted through the latent channels per data sample.

Original: Accordingly, we adopt an incremental form of the PI controller in Eq.
Revised: Accordingly, we adopt an incremental form of the PI controller in Eq.

Original: (8), and initialize it to a large value:
Revised: (8), and initialize it to a large value:

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: and β (0) is a large initial value.
Revised: and β (0) is a large initial value.

Original: When the PI controller is initialized to a large value β (0) , it can quickly produce a (small) KL-divergence during initial model training, preventing emergence of entangled factors.
Revised: When the PI controller is initialized to a large value β (0) , it can quickly produce a (small) KL-divergence during initial model training, preventing emergence of entangled factors.

Original: [Figure]
Revised: [Figure]

Original: Moving average: Since our model is trained with mini-batch data, it often contains noise that causes β ( t ) to oscillate.
Revised: Moving average: Since our model is trained with mini-batch data, it often contains noise that causes β ( t ) to oscillate.

Original: In particular, when β ( t ) plunges during training, it would cause KL-divergence to rise too quickly.
Revised: In particular, when β ( t ) plunges during training, it would cause KL-divergence to rise too quickly.

Original: This may lead to multiple latent factors coming out together to be entangled.
Revised: This may lead to multiple latent factors coming out together to be entangled.

Original: To mitigate this issue, we adopt moving average to smooth the output KL-divergence as the feedback of
Revised: To mitigate this issue, we adopt moving average to smooth the output KL-divergence as the feedback of

Original: PI controller below.
Revised: PI controller below.

Original: [Equation]
Revised: [Equation]

Original: where α i denotes weight and T denotes the window size of past training steps.
Revised: where α i denotes weight and T denotes the window size of past training steps.

Original: Hybrid annealing: Control systems with step (input) function (i.e., those where the set point can change abruptly) often suffer from an overshoot problem.
Revised: Hybrid annealing: Control systems with step (input) function (i.e., those where the set point can change abruptly) often suffer from an overshoot problem.

Original: An overshoot is temporary overcompensation, where the controlled variable oscillates around the set point.
Revised: An overshoot is temporary overcompensation, where the controlled variable oscillates around the set point.

Original: In our case, it means that the actual
Revised: In our case, it means that the actual

Original: KL-divergence may signiﬁcantly (albeit temporarily) exceed the desired value, when set point is abruptly changed.
Revised: KL-divergence may signiﬁcantly (albeit temporarily) exceed the desired value, when set point is abruptly changed.

Original: This effect would cause some latent factors to come out earlier than expected, and be entangled, thereby producing poor-quality disentanglement.
Revised: This effect would cause some latent factors to come out earlier than expected, and be entangled, thereby producing poor-quality disentanglement.

Original: To address this problem, we develop a hybrid annealing method that changes the set point more gradually, as illustrated in Fig.
Revised: To address this problem, we develop a hybrid annealing method that changes the set point more gradually, as illustrated in Fig.

Original: 7 in Appendix.
Revised: 7 in Appendix.

Original: It combines step function with ramp function to smoothly increase the target KL-divergence in order to prevent overshoot and thus better disentangle latent factors one by one.
Revised: It combines step function with ramp function to smoothly increase the target KL-divergence in order to prevent overshoot and thus better disentangle latent factors one by one.

Original: The combination of the above three methods allows DynamicVAE to evolve β ( t ) along a favorable trajectory to separate disentanglement learning and reconstruction optimization.
Revised: The combination of the above three methods allows DynamicVAE to evolve β ( t ) along a favorable trajectory to separate disentanglement learning and reconstruction optimization.

Original: We summarize the proposed incremental PI algorithm in Algorithm 1 in Appendix B.
Revised: We summarize the proposed incremental PI algorithm in Algorithm 1 in Appendix B.

Original: 3. S TABILITY A NALYSIS OF D YNAMIC VAE
Revised: 3. S TABILITY A NALYSIS OF D YNAMIC VAE

Original: We further theoretically analyze the stability of the proposed DynamicVAE.
Revised: We further theoretically analyze the stability of the proposed DynamicVAE.

Original: This work is the ﬁrst to offer the necessary and sufﬁcient conditions that control hyperparameters should satisfy in order to guarantee the stability of KL-divergence, when β is manipulated dynamically during the training process of a (variant of) β -VAE.
Revised: This work is the ﬁrst to offer the necessary and sufﬁcient conditions that control hyperparameters should satisfy in order to guarantee the stability of KL-divergence, when β is manipulated dynamically during the training process of a (variant of) β -VAE.

Original: To this end, our ﬁrst step is to build the state space model for our control system.
Revised: To this end, our ﬁrst step is to build the state space model for our control system.

Original: Throughout the paper, the state variable at training step t is deﬁned as x ( t ) =
Revised: Throughout the paper, the state variable at training step t is deﬁned as x ( t ) =

Original: β ( t ) .
Revised: β ( t ) .

Original: Accordingly, the model of incremental PI controller can be written as:
Revised: Accordingly, the model of incremental PI controller can be written as:

Original: [Equation]
Revised: [Equation]

Original: where error e ( t ) , as shown in Fig.
Revised: where error e ( t ) , as shown in Fig.

Original: 1(a), is given by e ( t ) =
Revised: 1(a), is given by e ( t ) =

Original: C − y ( t − 1) .
Revised: C − y ( t − 1) .

Original: Here y ( t ) is a dynamic model about the time response of the output KL divergence, y KL ( t ) .
Revised: Here y ( t ) is a dynamic model about the time response of the output KL divergence, y KL ( t ) .

Original: According to Liu & Theodorou (2019), stochastic gradient descent (SGD) for optimizing an objective function can be described by a ﬁrst-order dynamic model.
Revised: According to Liu & Theodorou (2019), stochastic gradient descent (SGD) for optimizing an objective function can be described by a ﬁrst-order dynamic model.

Original: Our experiment, as illustrated in Fig. 1(b), also shows that the response y ( t ) in the open loop system approximately meets a negative exponential function, which further veriﬁes that our system is a ﬁrst-order dynamic system.
Revised: Our experiment, as illustrated in Fig. 1(b), also shows that the response y ( t ) in the open loop system approximately meets a negative exponential function, which further veriﬁes that our system is a ﬁrst-order dynamic system.

Original: We hence use the ﬁrst-order dynamic model to describe it below.
Revised: We hence use the ﬁrst-order dynamic model to describe it below.

Original: [Equation]
Revised: [Equation]

Original: (13) where a is a positive hyperparameter to describe the dynamic property, and g ( x ) is a mapping function between the actual KL-divergence and β ( t ) .
Revised: (13) where a is a positive hyperparameter to describe the dynamic property, and g ( x ) is a mapping function between the actual KL-divergence and β ( t ) .

Original: Since DynamicVAE is a discrete control system with sampling period T s = 1 , the above ﬁrst-order dynamic model can be reformulated as
Revised: Since DynamicVAE is a discrete control system with sampling period T s = 1 , the above ﬁrst-order dynamic model can be reformulated as

Original: [Equation]
Revised: [Equation]

Original: Now let x 1 ( t ) =
Revised: Now let x 1 ( t ) =

Original: x ( t ) , x 2 ( t ) =
Revised: x ( t ) , x 2 ( t ) =

Original: y ( t − 1) , x 3 ( t ) =
Revised: y ( t − 1) , x 3 ( t ) =

Original: y ( t − 2) , then Eqs.
Revised: y ( t − 2) , then Eqs.

Original: (12) and (14) can be rewritten as the following state space equations.
Revised: (12) and (14) can be rewritten as the following state space equations.

Original: [Equation]
Revised: [Equation]

Original: In order to analyze the stability of the above non-linear state space model, one commonly used method is to linearize it at an equilibrium point (Hughes, 2015).
Revised: In order to analyze the stability of the above non-linear state space model, one commonly used method is to linearize it at an equilibrium point (Hughes, 2015).

Original: In this paper, we use the following equilibrium point:
Revised: In this paper, we use the following equilibrium point:

Original: [Equation]
Revised: [Equation]

Original: where g − 1 ( · ) denotes the inverse function and x ∗ 2 = x ∗ 3 .
Revised: where g − 1 ( · ) denotes the inverse function and x ∗ 2 = x ∗ 3 .

Original: Next, we apply the ﬁrst-order Taylor expansion to the above Eq.
Revised: Next, we apply the ﬁrst-order Taylor expansion to the above Eq.

Original: (15), yielding
Revised: (15), yielding

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: and A is the Jacobian matrix at equilibrium point x ∗ , as deﬁned in Eq.
Revised: and A is the Jacobian matrix at equilibrium point x ∗ , as deﬁned in Eq.

Original: (20) in Appendix C. After this linearization, we can prove the stability of the proposed method as the modulus of eigenvalue λ of A is smaller than 1 , as described in the following theorem.
Revised: (20) in Appendix C. After this linearization, we can prove the stability of the proposed method as the modulus of eigenvalue λ of A is smaller than 1 , as described in the following theorem.

Original: Theorem 3.1.
Revised: Theorem 3.1.

Original: Let a > 0 and assume g (cid:48)
Revised: Let a > 0 and assume g (cid:48)

Original: ( x ) < 0 , ∀ x > 0 .
Revised: ( x ) < 0 , ∀ x > 0 .

Original: Then DynamicVAE is stable at the equilibrium point C if and only if the parameters of the PI controller, K i and K p , satisfy the following conditions
Revised: Then DynamicVAE is stable at the equilibrium point C if and only if the parameters of the PI controller, K i and K p , satisfy the following conditions

Original: [Equation]
Revised: [Equation]

Original: We provide the detailed proof in Appendix C.
Revised: We provide the detailed proof in Appendix C.

Original: Remark 3.1.
Revised: Remark 3.1.

Original: The assumption of g (cid:48) ( x ) < 0 , ∀ x basically asks that the KL term in the objective to be a monotonously decreasing function of the coefﬁcient β ( t ) , and we also further empirically corroborate its validity on two benchmark datasets as shown in Appendix C.1.
Revised: The assumption of g (cid:48) ( x ) < 0 , ∀ x basically asks that the KL term in the objective to be a monotonously decreasing function of the coefﬁcient β ( t ) , and we also further empirically corroborate its validity on two benchmark datasets as shown in Appendix C.1.

Original: In addition, we choose K p and K i that meet the above conditions (19) to verify the stability of DynamicVAE in Appendix C.1.
Revised: In addition, we choose K p and K i that meet the above conditions (19) to verify the stability of DynamicVAE in Appendix C.1.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In this section we evaluate the performance of DynamicVAE and compare it against existing baselines, including ControlVAE (Shao et al., 2020), β -VAE H (Higgins et al., 2017b), β -VAE B (Burgess et al., 2018), FactorVAE (Kim & Mnih, 2018), and VAE (Kingma & Welling, 2013).
Revised: In this section we evaluate the performance of DynamicVAE and compare it against existing baselines, including ControlVAE (Shao et al., 2020), β -VAE H (Higgins et al., 2017b), β -VAE B (Burgess et al., 2018), FactorVAE (Kim & Mnih, 2018), and VAE (Kingma & Welling, 2013).

Original: We conduct experiments on three benchmark datasets: dSprites (Burgess et al., 2018), MNIST (Chen et al., 2016) and 3D Chairs (Aubry et al., 2014).
Revised: We conduct experiments on three benchmark datasets: dSprites (Burgess et al., 2018), MNIST (Chen et al., 2016) and 3D Chairs (Aubry et al., 2014).

Original: The detailed model conﬁgurations and hyperparameter settings are presented in Appendix A. Source code will be publicly available upon publication.
Revised: The detailed model conﬁgurations and hyperparameter settings are presented in Appendix A. Source code will be publicly available upon publication.

Original: 4. R ESULTS AND A NALYSIS
Revised: 4. R ESULTS AND A NALYSIS

Original: Dsprites Dataset: We ﬁrst evaluate the performance of DynamicVAE on learning disentangled representations using dSprites .
Revised: Dsprites Dataset: We ﬁrst evaluate the performance of DynamicVAE on learning disentangled representations using dSprites .

Original: Fig.
Revised: Fig.

Original: 2 (a) and (b) illustrate the comparison of reconstruction error and the hyperparameter β ( t ) (using 5 random seeds) for different approaches.
Revised: 2 (a) and (b) illustrate the comparison of reconstruction error and the hyperparameter β ( t ) (using 5 random seeds) for different approaches.

Original: We can observe from
Revised: We can observe from

Original: Fig.
Revised: Fig.

Original: 2 (a) that DynamicVAE (KL=20) has much lower reconstruction error (about 11 . 8 ) than β -VAE and FactorVAE, and is comparable to the basic VAE and ControlVAE.
Revised: 2 (a) that DynamicVAE (KL=20) has much lower reconstruction error (about 11 . 8 ) than β -VAE and FactorVAE, and is comparable to the basic VAE and ControlVAE.

Original: This is because DynamicVAE dynamically adjusts the weight, β ( t ) , to balance the disentanglement and reconstruction.
Revised: This is because DynamicVAE dynamically adjusts the weight, β ( t ) , to balance the disentanglement and reconstruction.

Original: Speciﬁcally,
Revised: Speciﬁcally,

Original: [Figure]
Revised: [Figure]

Original: = errors
Revised: = errors

Original: DynamicVAE turns the
Revised: DynamicVAE turns the

Original: [Table]
Revised: [Table]

Original: DynamicVAE automatically assigns a large β ( t ) at the beginning of training in order to obtain good disentanglement, and then its weight gradually drops to less than 1 at the end of optimization, as shown in Fig.
Revised: DynamicVAE automatically assigns a large β ( t ) at the beginning of training in order to obtain good disentanglement, and then its weight gradually drops to less than 1 at the end of optimization, as shown in Fig.

Original: 2 (b).
Revised: 2 (b).

Original: In contrast, β -VAE and FactorVAE have a large and ﬁxed weight in the objective so that their optimization algorithms tend to optimize the KL-divergence term (total correlation term for FactorVAE), leading to higher reconstruction error.
Revised: In contrast, β -VAE and FactorVAE have a large and ﬁxed weight in the objective so that their optimization algorithms tend to optimize the KL-divergence term (total correlation term for FactorVAE), leading to higher reconstruction error.

Original: For ControlVAE, it can also dynamically tune β ( t ) to control the value of KL-divergence, but its disentanglement performance degrades with the increase of KL-divergence (i.e., decrease of weight) as illustrated in Table 1.
Revised: For ControlVAE, it can also dynamically tune β ( t ) to control the value of KL-divergence, but its disentanglement performance degrades with the increase of KL-divergence (i.e., decrease of weight) as illustrated in Table 1.

Original: In addition, Fig. 2(c) illustrates an example of per-factor KL-divergence in the latent code as the total information capacity (KL-divergence) increases from 0 .
Revised: In addition, Fig. 2(c) illustrates an example of per-factor KL-divergence in the latent code as the total information capacity (KL-divergence) increases from 0 .

Original: 5 to 20 .
Revised: 5 to 20 .

Original: We can see that DynamicVAE disentangles all the ﬁve data generative factors, starting from position ( x and y ) to scale, followed by orientation and shape.
Revised: We can see that DynamicVAE disentangles all the ﬁve data generative factors, starting from position ( x and y ) to scale, followed by orientation and shape.

Original: Next, we use a robust disentanglement metric, robust mutual information gap (RMIG) (Do &
Revised: Next, we use a robust disentanglement metric, robust mutual information gap (RMIG) (Do &

Original: Tran, 2020), to evaluate the disentanglement of different methods.
Revised: Tran, 2020), to evaluate the disentanglement of different methods.

Original: We can see from Table 1 that
Revised: We can see from Table 1 that

Original: DynamicVAE has a comparable RMIG score to the FactorVAE, but it has much lower reconstruction error as illustrated in Fig.
Revised: DynamicVAE has a comparable RMIG score to the FactorVAE, but it has much lower reconstruction error as illustrated in Fig.

Original: Moreover, DynamicVAE has higher RMIG score but lower reconstruction error than β -VAE models.
Revised: Moreover, DynamicVAE has higher RMIG score but lower reconstruction error than β -VAE models.

Original: We also ﬁnd that our method achieves much better disentanglement than
Revised: We also ﬁnd that our method achieves much better disentanglement than

Original: ControlVAE for comparable reconstruction accuracy.
Revised: ControlVAE for comparable reconstruction accuracy.

Original: Hence, DynamicVAE is able to improve the reconstruction quality yet obtain good disentanglement.
Revised: Hence, DynamicVAE is able to improve the reconstruction quality yet obtain good disentanglement.

Original: Qualitatively, we also visualize the disentanglement results of different models in Fig.
Revised: Qualitatively, we also visualize the disentanglement results of different models in Fig.

Original: We can observe that DynamicVAE disentangles all the ﬁve generative factors on dSprites.
Revised: We can observe that DynamicVAE disentangles all the ﬁve generative factors on dSprites.

Original: However,
Revised: However,

Original: ControlVAE is not very effective to disentangle all the factors when its KL-divergence is set to a large value, such as 20 .
Revised: ControlVAE is not very effective to disentangle all the factors when its KL-divergence is set to a large value, such as 20 .

Original: Furthermore, β -VAE B ( γ = 100 ) disentangles four generative factors and mistakenly combines the scale and shape factors together (in the third row).
Revised: Furthermore, β -VAE B ( γ = 100 ) disentangles four generative factors and mistakenly combines the scale and shape factors together (in the third row).

Original: The other methods do not perform well for disentanglement.
Revised: The other methods do not perform well for disentanglement.

Original: MNIST and 3D Chair Datasets
Revised: MNIST and 3D Chair Datasets

Original: We also evaluate the proposed method on the other two datasets:
Revised: We also evaluate the proposed method on the other two datasets:

Original: MNIST and 3D Chairs.
Revised: MNIST and 3D Chairs.

Original: Fig.
Revised: Fig.

Original: 4 illustrates some samples of the disentangled factors for DynamicVAE on MNIST.
Revised: 4 illustrates some samples of the disentangled factors for DynamicVAE on MNIST.

Original: We also verify that our method achieves better disentanglement compared with the other methods, shown in Appendix D.
Revised: We also verify that our method achieves better disentanglement compared with the other methods, shown in Appendix D.

Original: In addition, our method with β < 1 can signiﬁcantly improve the reconstruction accuracy than β -VAE as illustrated in Fig.
Revised: In addition, our method with β < 1 can signiﬁcantly improve the reconstruction accuracy than β -VAE as illustrated in Fig.

Original: 10 in Appendix D.
Revised: 10 in Appendix D.

Original: We also demonstrate that DynamicVAE can learn many different data generative factors on another challenging dataset, 3D Chairs.
Revised: We also demonstrate that DynamicVAE can learn many different data generative factors on another challenging dataset, 3D Chairs.

Original: We can observe from Fig.
Revised: We can observe from Fig.

Original: 5 that our method disentangles six different latent factors, such as wheels, and leg height and azimuth, same as FactorVAE (Kim & Mnih, 2018).
Revised: 5 that our method disentangles six different latent factors, such as wheels, and leg height and azimuth, same as FactorVAE (Kim & Mnih, 2018).

Original: ControlVAE (KL=20) !
Revised: ControlVAE (KL=20) !

Original: -VAE B ( ! = 100 ) x
Revised: -VAE B ( ! = 100 ) x

Original: y sc a l e o r i en t s hape
Revised: y sc a l e o r i en t s hape

Original: FactorVAE ( !
Revised: FactorVAE ( !

Original: = 10 ) x y sc a l e o r i en t s hape
Revised: = 10 ) x y sc a l e o r i en t s hape

Original: DynamicVAE (KL=20)
Revised: DynamicVAE (KL=20)

Original: Basic VAE !
Revised: Basic VAE !

Original: -VAE H
Revised: -VAE H

Original: ( ! = 4 )
Revised: ( ! = 4 )

Original: [Figure]
Revised: [Figure]

Original: Decoupled Reconstruction and Disentanglement
Revised: Decoupled Reconstruction and Disentanglement

Original: Additionally, we show that the proposed DynamicVAE is able to decouple the reconstruction and disentanglement learning into two phases, overcoming the problem of balancing the trade-off between reconstruction and disentanglement.
Revised: Additionally, we show that the proposed DynamicVAE is able to decouple the reconstruction and disentanglement learning into two phases, overcoming the problem of balancing the trade-off between reconstruction and disentanglement.

Original: Fig.
Revised: Fig.

Original: 6 illustrates the RMIG score and reconstruction loss with the increase of training steps after all the factors are disentangled (before 800 , 000 ).
Revised: 6 illustrates the RMIG score and reconstruction loss with the increase of training steps after all the factors are disentangled (before 800 , 000 ).

Original: It can be seen that RMIG score of our method remains stable as the reconstruction loss drops.
Revised: It can be seen that RMIG score of our method remains stable as the reconstruction loss drops.

Original: Therefore, the proposed method does not introduce any conﬂict between reconstruction  and disentanglement learning
Revised: Therefore, the proposed method does not introduce any conﬂict between reconstruction optimization and disentangled representation learning

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original:  with training steps.
Revised: vary with training steps.

Original: 4. A BLATION S TUDIES
Revised: 4. A BLATION S TUDIES

Original: To compare the performance of DynamicVAE and its variants, we perform following ablation studies:
Revised: To compare the performance of DynamicVAE and its variants, we perform following ablation studies:

Original: • DynamicVAE-P: it uses positional PI controller with no initialization to a large β (0) , instead of the incremental PI initialized to a large β (0) , to tune the weight on KL term in the VAE objective.
Revised: • DynamicVAE-P: it uses positional PI controller with no initialization to a large β (0) , instead of the incremental PI initialized to a large β (0) , to tune the weight on KL term in the VAE objective.

Original: • DynamicVAE-step: it solely adopts step function without ramp function for our annealing method.
Revised: • DynamicVAE-step: it solely adopts step function without ramp function for our annealing method.

Original: • DynamicVAE-t: this model directly uses the output KL-divergence at time t as a feedback of PI controller without using moving average to smooth it.
Revised: • DynamicVAE-t: this model directly uses the output KL-divergence at time t as a feedback of PI controller without using moving average to smooth it.

Original: [Table]
Revised: [Table]

Original: Table 2 shows the comparison of RMIG score for DynamicVAE and its variants.
Revised: Table 2 shows the comparison of RMIG score for DynamicVAE and its variants.

Original: It can be observed that DynamicVAE outperforms the other methods in terms of overall RMIG score.
Revised: It can be observed that DynamicVAE outperforms the other methods in terms of overall RMIG score.

Original: We also ﬁnd
Revised: We also ﬁnd

Original: that
Revised: that

Original: DynamicVAE-step does not perform well because the ramp function is removed from our annealing method, leading to overshoot of PI controller.
Revised: DynamicVAE-step does not perform well because the ramp function is removed from our annealing method, leading to overshoot of PI controller.

Original: As a result, it makes the other factors come out earlier and entangled to each other.
Revised: As a result, it makes the other factors come out earlier and entangled to each other.

Original: Thus, we can conclude the importance of adding ramp function for our annealing method.
Revised: Thus, we can conclude the importance of adding ramp function for our annealing method.

Original: In addition, we can see that the proposed moving average and incremental PI control algorithm also play a critical role to improve the disentanglement.
Revised: In addition, we can see that the proposed moving average and incremental PI control algorithm also play a critical role to improve the disentanglement.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Disentangled representation learning can be divided into two main categories: unsupervised learning and supervised learning.
Revised: Disentangled representation learning can be divided into two main categories: unsupervised learning and supervised learning.

Original: Supervised disentanglement learning (Mathieu et al., 2016; Siddharth et al., 2017; Kingma et al., 2014;
Revised: Supervised disentanglement learning (Mathieu et al., 2016; Siddharth et al., 2017; Kingma et al., 2014;

Original: Reed et al., 2014) requires the prior knowledge of some data generative factors from human annotation to train the model.
Revised: Reed et al., 2014) requires the prior knowledge of some data generative factors from human annotation to train the model.

Original: Some studies (Locatello et al., 2019a;b;c) ﬁgure out that it is hard to achieve reliable and good disentanglement without supervision.
Revised: Some studies (Locatello et al., 2019a;b;c) ﬁgure out that it is hard to achieve reliable and good disentanglement without supervision.

Original: For supervised learning, the limited labeling information can help ensure a latent space of the VAE with desirable structure w.r.t to the ground-truth latent factors.
Revised: For supervised learning, the limited labeling information can help ensure a latent space of the VAE with desirable structure w.r.t to the ground-truth latent factors.

Original: In order to reduce human annotations, researchers tried to develop weakly supervised learning (Bouchacourt et al., 2018; Hosoya, 2019; Locatello et al., 2020) to learn disentangled representations.
Revised: In order to reduce human annotations, researchers tried to develop weakly supervised learning (Bouchacourt et al., 2018; Hosoya, 2019; Locatello et al., 2020) to learn disentangled representations.

Original: However, these methods still require explicit human labeling or assume the change of the two observations is small.
Revised: However, these methods still require explicit human labeling or assume the change of the two observations is small.

Original: In practice, it is unrealistic for initial learners to discover the data generative factors in most real world scenarios.
Revised: In practice, it is unrealistic for initial learners to discover the data generative factors in most real world scenarios.

Original: For unsupervised learning methods, the recent approaches mainly build on Variational Autoencoders (VAEs) (Kingma & Welling, 2013) and Generative Adversarial Networks (GANs) (Goodfellow et al., 2014).
Revised: For unsupervised learning methods, the recent approaches mainly build on Variational Autoencoders (VAEs) (Kingma & Welling, 2013) and Generative Adversarial Networks (GANs) (Goodfellow et al., 2014).

Original: InfoGAN (Lin et al., 2020) is the ﬁrst scalable unsupervised learning method for disentangling.
Revised: InfoGAN (Lin et al., 2020) is the ﬁrst scalable unsupervised learning method for disentangling.

Original: It, however, suffers from training instabilities and does not perform well in disentanglement learning (Higgins et al., 2018), so most recent works are largely based on VAEs models.
Revised: It, however, suffers from training instabilities and does not perform well in disentanglement learning (Higgins et al., 2018), so most recent works are largely based on VAEs models.

Original: The VAE models, such as β -VAE ( β > 1 ), FactorVAE and β -TCVAE
Revised: The VAE models, such as β -VAE ( β > 1 ), FactorVAE and β -TCVAE

Original: (Chen et al., 2018), often suffer from high reconstruction errors in order to obtain better disentanglement, since they add a large weight to terms in the objective.
Revised: (Chen et al., 2018), often suffer from high reconstruction errors in order to obtain better disentanglement, since they add a large weight to terms in the objective.

Original: To address this issue, Shao et al.
Revised: To address this issue, Shao et al.

Original: (2020) develops a controllable variational autoencoder (ControlVAE) to dynamically tune β to achieve the trade-off between reconstruction quality and disentanglement.
Revised: (2020) develops a controllable variational autoencoder (ControlVAE) to dynamically tune β to achieve the trade-off between reconstruction quality and disentanglement.

Original: However, ControlVAE fails to decouple disentanglementlearning and reconstruction  because the computed β may oscillate during training, since it uses step function as annealing method and frequently adjusts β at each step.
Revised: However, ControlVAE fails to decouple disentanglement and reconstruction accuracy because the computed β may oscillate during training, since it uses step function as annealing method and frequently adjusts β at each step.

Original: In this paper, we propose a new method,
Revised: In this paper, we propose a new method,

Original: DynamicVAE, that can separate disentanglement learning and reconstruction.
Revised: DynamicVAE, that can separate disentangled representation learning and reconstruction.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: This paper proposed a novel dynamic learning method, DynamicVAE, to address the trade-off problem between reconstruction and disentanglement.
Revised: This paper proposed a novel dynamic learning method, DynamicVAE, to address the trade-off problem between reconstruction and disentanglement.

Original: Our method is able to turn the weight of β -VAE to a small value ( β ≤ 1 ) to achieve good disentanglement against the previous default consumption, β > 1 .
Revised: Our method is able to turn the weight of β -VAE to a small value ( β ≤ 1 ) to achieve good disentanglement against the previous default consumption, β > 1 .

Original: Speciﬁcally, we leveraged an incremental PI controller, moving average and a hybrid annealing to stabilize the KL-divergence to decouple the disentanglement and reconstruction.
Revised: Speciﬁcally, we leveraged an incremental PI controller, moving average and a hybrid annealing to stabilize the KL-divergence to decouple the disentanglement and reconstruction.

Original: We further theoretically prove the stability of the proposed DynamicVAE.
Revised: We further theoretically prove the stability of the proposed DynamicVAE.

Original: The evaluation results demonstrate DynamicVAE can signiﬁcantly improve the reconstruction accuracy meanwhile attaining good disentanglement. 
Revised: The evaluation results demonstrate DynamicVAE can signiﬁcantly improve the reconstruction accuracy meanwhile attaining good disentanglement. It decouples disentanglement and reconstruction accuracy without introducing a conﬂict between them.

Revised: A M ODEL C ONFIGURATIONS AND H YPERPARAMETER S ETTINGS

Revised: We summarize the detailed model conﬁgurations and hyperparameter settings for DynamicVAE below.

Revised: Following the same model architecture of β -VAE, we adopt a convolutional layer and deconvolutional layer for our experiments.

Revised: We use Adam optimizer with β 1 = 0 .

Original: It decouples disentanglement learning andreconstruction without introducing a conﬂict between them.
Revised: 90 , β 2 = 0 . 99 and a learning rate tuned from 10 − 4 . We set K p and K i for PI algorithm to 0 . 01 and 0 .

Revised: 005 , respectively.

Revised: The weight β ( t ) for incremental PI controller is initialized with 150 , 100 and 50 for dSprites, MNIST and 3D Chairs, respectively.

Revised: The batch size is set to 128 .

Revised: Using the similar methodology in (Burgess et al., 2018), we train a single model by gradually increasing KL-divergence from 0 .

Revised: 5 to a desired value C with a step function s and ramp function for every M training steps, as shown in Fig.

Revised: In the experiment, we set the step, s , to 0 .

Revised: 15 per M = 6 , 000 training steps (including 5 , 000 in step function and 1 ,in ramp function) as the information capacity (desired KL- divergence) increases from 0 .

Revised: 5 until to 20 , 26 and 18 for dSprites, MNIST and 3D Chairs datasets respectively.

Revised: In addition, the window size of moving average is T = 5 with equal weight α .

Revised: Our model adopts the same encoder and decoder architecture as β -VAE H and ControlVAE except for plugging in PI control algorithm, as illustrated in Table 3 and Table 4.

Revised: [Table]

Revised: [Table]

Revised: A. PI P ARAMETER T UNING AND S ET P OINT G UIDELINES

Revised: We can tune PI parameters by following the conditions that guarantee the stability of the proposed method in Eq.(19) in Appendix C.

Revised: In addition, β (0) is initialized to a sufﬁciently large value in order to guarantee the KL-divergence is closed to zero at the beginning of model training.

Revised: On the other hand, the choice of desired value of KL-divergence (set point) is very simple.

Revised: Since our method achieves good disentanglement when β = 1 , we can set its desired value to equal or lower than the

Revised: KL-divergence of the original VAE as it converges.

Revised: H YBRID A NNEALING M ETHOD

Revised: Fig.

Revised: 7 shows the hybrid annealing method that combines step function with ramp function to gradually increase the desired value of KL-divergence for DynamicVAE.

Revised: B A LGORITHM

Revised: We summarize the incremental PI algorithm in Algorithm 1.

Revised: C P ROOF OF S TABILITY IN T HEOREM 3.

Revised: In this section we provide the omitted proof in the main paper about Theorem 3.1.

Revised: For convenience purpose, we ﬁrst restate Theorem 3.1 below.

Revised: Theorem 3.1.

Revised: Let a > 0 and assume g (cid:48)

Revised: ( x ) < 0 , ∀ x > 0 .

Revised: Then DynamicVAE is stable at the equilibrium point C if and only if the parameters of the PI controller, K i and K p , satisfy the following conditions

Revised: [Equation]

Revised: Proof.

Revised: At a high level

Revised: [Equation]

Revised: the proof goes by showing the spectral norm of the Jacobian matrix A to be strictly less than 1 under the given condition, which is both sufﬁcient and necessary for stability.

Revised: To start with, recall that the Jacobian matrix A at equilibrium point x ∗ is deﬁned by

Revised: [Equation]

Revised: where

Revised: [Equation]

Revised: ∂x

Revised: [Equation]

Revised: In order to guarantee the stability of our state space model, the modulus of eigenvalue λ of A should be smaller than 1, i.e., | λ | < 1 .

Revised: By deﬁnition, the eigenvalues of A can be obtained by computing the roots of the following characteristic polynomial:

Revised: [Equation]

Revised: Instead of computing the (complex) roots of the above cubic polynomial analytically, we use the following bilinear transformation (Jury, 1964) to map the unit circle | λ | < 1 to the left half plane such that its real root is less than 0 (Hughes, 2015):

Revised: [Equation]

Revised: Substituting λ in Eq.(22) with (23) , we have

Revised: [Equation]

Revised: where

Revised: [Equation]

Revised: Clearly, using the above transformation, we know that | λ | < 1 iff the real part of λ is less than 0, i.e.,

Revised: Re { ξ } < 0 .

Revised: In order to ensure Re { ξ } < 0 , based on Routh–Hurwitz stability criterion (Zabczyk,

Revised: [Figure]

Revised: 1992), b 0 , b 1 , b 2 , b 3 should satisfy the following sufﬁcient and necessary condition.

Revised: [Equation]

Revised: Substitute Eqs.

Revised: (21), (24) and (26) into the above system of inequalities, yielding  

Revised: [Equation]

Revised: To complete the proof, recall that a > 0 and we assume g (cid:48)

Original: {+( x ) < 
+}
Revised: {+( x ) < 
+}

Revised: , ∀ x .

Revised: Hence the coefﬁcients of PI controller, K p and K i in Eq.(27), need to satisfy the following conditions.

Revised: [Equation]

Revised: Since K p > 0 , K i > 0 in our designed PI control algorithm and g (cid:48) ( x ∗ 1 ) < 0 , we can further simplify it as

Revised: [Equation]

Revised: Therefore, as K p and K i meet the above conditions (29), our DynamicVAE would be stable at the set point, which is veriﬁed by the following experiments on different datasets.

Revised: (cid:4)

Revised: [Figure]

Revised: [Figure]

Revised: V ERIFICATION ON B ENCHMARK D ATASETS

Revised: We ﬁrst verify the validity of our assumption that g (cid:48) ( x ) < 0 in Theorem 3.1.

Revised: Fig.

Revised: 8 illustrates the relationship between β ( t ) and the actual KL when model training converges on dSprites and MNIST datasets.

Revised: We can observe that the actual output KL-divergence and β ( t ) have a highly negative correlation, which means g (cid:48)

Revised: ( x ) < 0 .

Revised: Next, we are going to verify the stability of the proposed DynamicVAE on MNIST and dSprites datasets.

Revised: On MNIST dataset, its mapping function g ( x ) in Fig.

Revised: 8 (a) can be approximately obtained by curve ﬁtting with the following negative exponential function:

Revised: [Equation]

Revised: And the corresponding derivative is

Revised: [Equation]

Revised: In addition, we introduce how to obtain the hyperparameter a in our dynamic model in Eq.

Revised: Assume that KL-divergence converges to a certain value C (cid:48) in the open loop control system during model training, then the dynamic model in Eq.

Revised: (14) can be rewritten as

Revised: [Equation]

Revised: When y (0) = 0 , and the sampling period of our system is T s = 1 , the corresponding solution is given by

Revised: [Equation]

Revised: In order to obtain the value of a , one commonly used method in control theory is to set a = 1 t ∗ as we have y ( t ∗ )

Revised: =

Revised: C (cid:48) (1 − exp( − 1))

Revised: ≈ 0 .

Revised: 632 C (cid:48) .

Revised: In this way, we can derive a based on the training steps t ∗ as KL-divergence reaches 63.2% of its ﬁnal value C (cid:48)

Revised: (Isa et al., 2011) in the experiments, as shown in Fig.

Revised: For MNIST dataset, we can get the hyperparameter a = 1 5000 around based on the time response of KL-divergence in the open loop system, as shown in Fig.

Revised: 9 (a).

Revised: Similarly, the derivative of mapping function on dSprites can be approximately expressed by

Revised: [Equation]

Revised: In addition, we can get the hyperparameter a = 1 2500 around based on the time response of KLdivergence in the open loop system, as shown in Fig.

Revised: 9 (b).

Revised: We summarize the parameters a and g (cid:48)

Revised: ( x ( t )) for different datasets in the following Table 5.

Revised: In this paper, we choose K p = 0 .

Revised: 01 and K i = 0 .

Revised: 005 with the parameters in Table 5 to validate our model meets the conditions in Eq.

Revised: In addition, our experimental results in Section 4 further demonstrate that our method can stabilize the KL-divergence to the set points.

Revised: [Figure]

Revised: [Table]

Revised: [Figure]

Revised: D E XTRA E XPERIMENTS ON MNIST D. E VALUATION ON R ECONSTRUCTION Q UALITY

Revised: Fig.

Revised: 10 shows the comparison of reconstruction loss and weight β ( t ) for different methods.

Revised: It can be observed that DynamicVAE and ControlVAE have comparable reconstruction accuracy to the basic

Revised: VAE, but they have better disentanglement than it, as shown in Fig.

Revised: 4 and 15.

Revised: In addition, we can see that DynamicVAE has better reconstruction quality than the two β -VAE models.

Revised: [Figure]

Revised: D. MNIST L ATENT T RAVERSALS FOR B ASELINES

Revised: We present some samples of latent traversals for the baseline methods.

Revised: We ﬁnd that DynamicVAE outperforms ControlVAE in term of rotation factor as illustrated in Fig.

Revised: 4 and 11, though they have comparable disentanglement score.

Revised: Moreover, DynamicVAE has a better disentanglement than β -VAE and FactorVAE in the following ﬁgures.

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]