SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yue Song $^{1,2}$ , T. Anderson Keller $^{2,3}$ , Nicu Sebe $^{1}$ , and Max Welling $^{2,3}$
$^{1}$ Department of Information Engineering and Computer Science, University of Trento, Italy  
 $^{2}$ Amsterdam Machine Learning Lab, University of Amsterdam, the Netherlands  
 $^{3}$ UvA-Bosch Delta Lab, University of Amsterdam, the Netherlands  
yue.song@unitn.it
A prominent goal of representation learning research is to achieve representations which are factorized in a useful manner with respect to the ground truth factors of variation. The fields of disentangled and equivariant representation learning have approached this ideal from a range of complimentary perspectives; however, to date, most approaches have proven to either be ill-specified or insufficiently flexible to effectively separate all realistic factors of interest in a learned latent space. In this work, we propose an alternative viewpoint on such structured representation learning which we call Flow Factorized Representation Learning, and demonstrate it to learn both more efficient and more usefully structured representations than existing frameworks. Specifically, we introduce a generative model which specifies a distinct set of latent probability paths that define different input transformations. Each latent flow is generated by the gradient field of a learned potential following dynamic optimal transport. Our novel setup brings new understandings to both disentanglement and equivariance. We show that our model achieves higher likelihoods on standard representation learning benchmarks while simultaneously being closer to approximately equivariant models. Furthermore, we demonstrate that the transformations learned by our model are flexibly composable and can also extrapolate to new data, implying a degree of robustness and generalizability approaching the ultimate goal of usefully factorized representation learning.
Developing models which learn useful representations of data has become an increasingly important focus in the machine learning community [5, 55]. For example, Large Language Models such as GPT [9] rely on an extensive pre-training phase to learn valuable representations, enabling quick finetuning on a diversity of tasks. However, a precise definition of what makes an ideal representation is still debated. One line of work has focused on 'disentanglement' of the underlying ground truth generative factors [5, 35, 13]. In general, the definition of 'disentanglement' often refers to learning and controlling statistically independent factors of variation [5, 36]. Over the years, many disentanglement methods have been proposed, including axis-aligned single-dimensional manipulation [35, 13], linear multi-dimensional traversals [78, 77, 90, 66], and, more recently, dynamic non-linear vector-based traversals [84, 79]. Although these methods have been met with significant success (and even linked to single-neuron brain activity [37, 91]), there are known theoretical limitations which make them ill-specified, including the presence of topological defects [7]. This has limited their deployment beyond toy settings.
Another line of work has focused on developing representations which respect symmetries of the underlying data in their output space [15, 36]. Specifically, equivariant representations are those for which the output transforms in a known predictable way for a given input transformation.
They can be seen to share many similarities with disentangled representations since an object undergoing a transformation which preserves its identity can be called a symmetry transformation [36]. Compared with disentanglement methods, equivariant networks are much more strictly defined, allowing for significantly greater control and theoretical guarantees with respect to the learned transformation [16, 50, 73, 20, 39]. However, this restriction also limits the types of transformations to which they may be applied. For example, currently only group transformations are supported, limiting real-world applicability. To avoid this caveat, some recent attempts propose to learn general but approximate equivariance from disentangled representations [49, 45, 79].
In this work, we provide an alternative viewpoint at the intersection of these two fields of work which we call Flow Factorized Representation Learning. Fig. 1 depicts the high-level illustration of our method. Given  $k$  different transformations  $p_k(\boldsymbol{x}_t|\boldsymbol{x}_0)$  in the input space, we have the corresponding latent probabilistic path  $\int_{\boldsymbol{z}_0,\boldsymbol{z}_t} q(\boldsymbol{z}_0|\boldsymbol{x}_0)q_k(\boldsymbol{z}_t|\boldsymbol{z}_0)p(\boldsymbol{x}_t|\boldsymbol{z}_t)$  for each of the transformations. Each latent flow path  $q_k(\boldsymbol{z}_t|\boldsymbol{z}_0)$  is generated by the gradient field of some learned potentials  $\nabla u^k$  following fluid mechanical dynamic Optimal Transport (OT) [4]. Our framework allows for novel understandings of both disentanglement and equivariance. The definition of disentanglement refers to the distinct set of tangent directions  $\nabla u^k$  that follow the OT paths to generate latent flows for modeling different factors of variation. The concept of equivariance in our case means that the two probabilistic paths, i.e.,  $p_k(\boldsymbol{x}_t|\boldsymbol{x}_0)$  in the image space and  $\int_{\boldsymbol{z}_0,\boldsymbol{z}_t} q(\boldsymbol{z}_0|\boldsymbol{x}_0)q_k(\boldsymbol{z}_t|\boldsymbol{z}_0)p(\boldsymbol{x}_t|\boldsymbol{z}_t)$  in the latent space, would eventually result in the same distribution of transformed data.
We build a formal generative model of sequences and integrate the above latent probability evolution as condition updates of the fac
torized sequence distribution. Based on the continuity equation, we derive a proper flow of probability density for the time evolution of both the prior and posterior. To perform inference, we approximate the true posterior of latent variables and train the parameters as a Variational Autoencoder (VAE) [47]. When the transformation type  $k$  is not observed (i.e., available as a label), we treat  $k$  as another latent variable and incorporate its posterior into our framework by learning it from sequences. Extensive experiments and thorough analyses have been conducted to show the effectiveness of our method. For example, we demonstrate empirically that our representations are usefully factorized, allowing flexible composability and generalization to new datasets. Furthermore, we show that our methods are also approximately equivariant by demonstrating that they commute with input transformations through the learned latent flows. Ultimately, we see these factors combine to yield the highest likelihood on the test set in each setting. Code is publicly available at https://github.com/KingJamesSong/latent-flow.
Figure 1: Illustration of our flow factorized representation learning: at each point in the latent space we have a distinct set of tangent directions  $\nabla u^k$  which define different transformations we would like to model in the image space. For each path, the latent sample evolves to the target on the potential landscape following dynamic optimal transport.
2 The generative model
In this section, we first introduce our generative model of sequences and then describe how we perform inference over the latent variables of this model in the next section.
2.1 Flow factorized sequence distributions
The model in this work defines a distribution over sequences of observed variables. We further factorize this distribution into  $k$  distinct components by assuming that each observed sequence is generated by one of the  $k$  separate flows of probability mass in latent space. Since in this work we
Figure 2: Depiction of our model in plate notation. (Left) Supervised, (Right) Weakly-supervised. White nodes denote latent variables, shaded nodes denote observed variables, solid lines denote the generative model, and dashed lines denote the approximate posterior. We see, as in a standard VAE framework, our model approximates the initial one-step posterior  $p(\pmb{z}_0|\pmb{x}_0)$ , but additionally approximates the conditional transition distribution  $p(\pmb{z}_t|\pmb{z}_{t-1}, k)$  through dynamic optimal transport over a potential landscape.
model discrete sequences of observations  $\bar{x} = \{x_0, x_1, \ldots, x_T\}$ , we aim to define a joint distribution with a similarly discrete sequence of latent variables  $\bar{z} = \{z_0, z_1, \ldots, z_T\}$ , and a categorical random variable  $k$  describing the sequence type (observed or unobserved). Explicitly, we assert the following factorization of the joint distribution over  $T$  timesteps:
$$
p (\bar {\boldsymbol {x}}, \bar {\boldsymbol {z}}, k) = p (k) p (\boldsymbol {z} _ {0}) p (\boldsymbol {x} _ {0} | \boldsymbol {z} _ {0}) \prod_ {t = 1} ^ {T} p (\boldsymbol {z} _ {t} | \boldsymbol {z} _ {t - 1}, k) p (\boldsymbol {x} _ {t} | \boldsymbol {z} _ {t}). \tag {1}
$$
Here  $p(k)$  is a categorical distribution defining the transformation type,  $p(\pmb{x}_t|\pmb{z}_t)$  asserts a mapping from latents to observations with Gaussian noise, and  $p(\pmb{z}_0) = \mathcal{N}(0,1)$ . A plate diagram of this model is depicted through the solid lines in Fig. 2.
2.2 Prior time evolution
To enforce that the time dynamics of the sequence define a proper flow of probability density, we compute the conditional update  $p(\pmb{z}_t|\pmb{z}_{t-1}, k)$  from the continuous form of the continuity equation:  $\partial_t p(\pmb{z}) = -\nabla \cdot (p(\pmb{z})\nabla \psi^k(\pmb{z}))$ , where  $\psi^k(\pmb{z})$  is the  $k$ 'th potential function which advects the density  $p(\pmb{z})$  through the induced velocity field  $\nabla \psi^k(\pmb{z})$ . Considering the discrete particle evolution corresponding to this density evolution,  $\pmb{z}_t = f(\pmb{z}_{t-1}, k) = \pmb{z}_{t-1} + \nabla_z \psi^k(\pmb{z}_{t-1})$ , we see that we can derive the conditional update from the continuous change of variables formula [69, 11]:
$$
p \left(\boldsymbol {z} _ {t} \mid \boldsymbol {z} _ {t - 1}, k\right) = p \left(\boldsymbol {z} _ {t - 1}\right) \left| \frac {d f \left(\boldsymbol {z} _ {t - 1} , k\right)}{d \boldsymbol {z} _ {t - 1}} \right| ^ {- 1} \tag {2}
$$
In this setting, we see that the choice of  $\psi$  ultimately determines the prior on the transition probability in our model. As a minimally informative prior for random trajectories, we use a diffusion equation achieved by simply taking  $\psi^k = -D_k\log p(\pmb{z}_t)$ . Then according to the continuity equation, the prior evolves as:
$$
\partial_ {t} p (\boldsymbol {z} _ {t}) = - \nabla \cdot \left(p (\boldsymbol {z} _ {t}) \nabla \psi\right) = D _ {k} \nabla^ {2} p (\boldsymbol {z} _ {t}) \tag {3}
$$
where  $D_{k}$  is a constant coefficient that does not change over time. The density evolution of the prior distribution thus follows a constant diffusion process. We set  $D_{k}$  as a learnable parameter which is distinct for each  $k$ .
3 Flow factorized variational autoencoders
To perform inference over the unobserved variables in our model, we propose to use a variational approximation to the true posterior, and train the parameters of the model as a VAE. To do this, we parameterize an approximate posterior for  $p(\pmb{z}_0|\pmb{x}_0)$ , and additionally parameterize a set of  $K$
functions  $u^{k}(z)$  to approximate the true latent potentials  $\psi^*$ . First, we will describe how we do this in the setting where the categorical random variable  $k$  is observed (which we call the supervised setting), then we will describe the model when  $k$  is also latent and thus additionally inferred (which we call the weakly supervised setting).
3.1 Inference with observed  $k$  (supervised)
When  $k$  is observed, we define our approximate posterior to factorize as follows:
$$
q \left(\bar {z} \mid \bar {x}, k\right) = q \left(z _ {0} \mid x _ {0}\right) \prod_ {t = 1} ^ {T} q \left(z _ {t} \mid z _ {t - 1}, k\right) \tag {4}
$$
We see that, in effect, our approximate posterior only considers information from element  $\mathbf{x}_0$ ; however, combined with supervision in the form of  $k$ , we find this is sufficient for the posterior to be able to accurately model full latent sequences. In the limitations section we discuss how the posterior could be changed to include all elements  $\{\mathbf{x}_t\}_{0}^{T}$  in future work.
Combing Eq. (4) with Eq. (1), we can derive the following lower bound to model evidence (ELBO):
$$
\begin{array}{l} \log p (\bar {\boldsymbol {x}} | k) = \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}}, k)} \left[ \log \frac {p (\bar {\boldsymbol {x}} , \bar {\boldsymbol {z}} | k)}{q (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}} , k)} \frac {q (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}} , k)}{p (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}} , k)} \right] \\ \geq \mathbb {E} _ {q _ {\theta} (\bar {z} | \bar {x}, k)} \left[ \log \frac {p (\bar {x} | \bar {z} , k) p (\bar {z} | k)}{q (\bar {z} | \bar {x} , k)} \right] \tag {5} \\ = \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}}, k)} \left[ \log p (\bar {\boldsymbol {x}} | \bar {\boldsymbol {z}}, k) \right] + \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}}, k)} \left[ \log \frac {p (\bar {\boldsymbol {z}} | k)}{q (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}} , k)} \right] \\ \end{array}
$$
Substituting and simplifying, Eq. (5) can be re-written as
$$
\begin{array}{l} \log p (\bar {\boldsymbol {x}} | k) \geq \sum_ {t = 0} ^ {T} \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}} | k)} \left[ \log p (\boldsymbol {x} _ {t} | \boldsymbol {z} _ {t}, k) \right] - \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}} | k)} \left[ \mathrm {D} _ {\mathrm {K L}} \left[ q _ {\theta} (\boldsymbol {z} _ {0} | \boldsymbol {x} _ {0}) | | p (\boldsymbol {z} _ {0}) \right] \right] \\ - \sum_ {t = 1} ^ {T} \mathbb {E} _ {q _ {\theta} (\bar {z} | k)} \left[ \mathrm {D} _ {\mathrm {K L}} \left[ q _ {\theta} \left(\boldsymbol {z} _ {t} \mid \boldsymbol {z} _ {t - 1}, k\right) \right| \left| p \left(\boldsymbol {z} _ {t} \mid \boldsymbol {z} _ {t - 1}, k\right) \right] \right] \\ \end{array}
$$
We thus see that we have an objective very similar to that of a traditional VAE, except that our posterior and our prior now both have a time evolution defined by the conditional distributions.
3.2 Inference with latent  $k$  (weakly supervised)
When  $k$  is not observed, we can treat it as another latent variable, and simultaneously perform inference over it in addition to the sequential latent  $\bar{z}$ . To achieve this, we define our approximate posterior and instead factorize it as
$$
q (\bar {z}, k | \bar {x}) = q (k | \bar {x}) q \left(\boldsymbol {z} _ {0} | \boldsymbol {x} _ {0}\right) \prod_ {t = 1} ^ {T} q \left(\boldsymbol {z} _ {t} | \boldsymbol {z} _ {t - 1}, k\right) \tag {7}
$$
Following a similar procedure as in the supervised setting, we derive the new ELBO as
$$
\begin{array}{l} \log p (\bar {\boldsymbol {x}}) = \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}}, k | \bar {\boldsymbol {x}})} \left[ \log \frac {p (\bar {\boldsymbol {x}} , \bar {\boldsymbol {z}} , k)}{q (\bar {\boldsymbol {z}} , k | \bar {\boldsymbol {x}})} \frac {q (\bar {\boldsymbol {z}} , k | \bar {\boldsymbol {x}})}{p (\bar {\boldsymbol {z}} , k | \bar {\boldsymbol {x}})} \right] \\ \geq \mathbb {E} _ {q _ {\theta} (\bar {z}, k | \bar {x})} \left[ \log \frac {p (\bar {x} | \bar {z} , k) p (\bar {z} | k)}{q (\bar {z} | \bar {x} , k)} \frac {p (k)}{q (k | \bar {x})} \right] \tag {8} \\ = \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}}, k | \bar {\boldsymbol {x}})} \left[ \log p (\bar {\boldsymbol {x}} | \bar {\boldsymbol {z}}, k) \right] + \mathbb {E} _ {q _ {\theta} (\bar {\boldsymbol {z}}, k | \bar {\boldsymbol {x}})} \left[ \log \frac {p (\bar {\boldsymbol {z}} | k)}{q (\bar {\boldsymbol {z}} | \bar {\boldsymbol {x}} , k)} \right] + \mathbb {E} _ {q _ {\gamma} (k | \bar {\boldsymbol {x}})} \left[ \log \frac {p (k)}{q (k | \bar {\boldsymbol {x}})} \right] \\ \end{array}
$$
We see that, compared with Eq. (5), only one additional KL divergence term  $\mathrm{D}_{\mathrm{KL}}[q_{\gamma}(k|\bar{\boldsymbol{x}})||p(k)]$  is added. The prior  $p(k)$  is set to follow a categorical distribution, and we apply the Gumbel-SoftMax trick [43] to allow for categorical re-parameterization and sampling of  $q_{\gamma}(k|\bar{\boldsymbol{x}})$ .
3.3 Posterior time evolution
As noted, to approximate the true generative model which has some unknown latent potentials  $\psi^k$ , we propose to parameterize a set of potentials as  $u^k(z,t) = \mathrm{MLP}([z;t])$  and train them through the ELBOs above. Again, we use the continuity equation to define the time evolution of the posterior, and thus we can derive the conditional time update  $q(\pmb{z}_t|\pmb{z}_{t-1},k)$  through the change of variables formula. Given the function of the sample evolution  $\pmb{z}_t = g(\pmb{z}_{t-1},k) = \pmb{z}_{t-1} + \nabla_z u^k$ , we have:
$$
q \left(\boldsymbol {z} _ {t} \mid \boldsymbol {z} _ {t - 1}, k\right) = q \left(\boldsymbol {z} _ {t - 1}\right) \left| \frac {d g \left(\boldsymbol {z} _ {t - 1} , k\right)}{d \boldsymbol {z} _ {t - 1}} \right| ^ {- 1} \tag {9}
$$
Converting the above continuous equation to the discrete setting and taking the logarithm of both sides gives the normalizing-flow-like density evolution of our posterior:
$$
\log q \left(\boldsymbol {z} _ {t} \mid \boldsymbol {z} _ {t - 1}, k\right) = \log q \left(\boldsymbol {z} _ {t - 1}\right) - \log \left| 1 + \nabla_ {\boldsymbol {z}} ^ {2} u ^ {k} \right| \tag {10}
$$
The above relation can be equivalently derived from the continuity equation (i.e.,  $\partial_t q(z) = -\nabla \cdot (q(z)\nabla u^k)$ ). Notice that we only assume the initial posterior  $q(z_0|\pmb{x}_0)$  follows a Gaussian distribution. For future timesteps, we do not pose any further assumptions and just let the density evolve according to the sample motion.
3.4 Ensuring optimal transport of the posterior flow
As an inductive bias, we would like each latent posterior flow to follow the OT path. To accomplish this, it is known that when the gradient  $\nabla u^k$  satisfies certain PDEs, the evolution of the probability density can be seen to minimize the  $L_{2}$  Wasserstein distance between the source distribution and the distribution of the target transformation. Specifically, we have:
Theorem 1 (Benamou-Brenier Formula [4]). For probability measures  $\mu_0$  and  $\mu_{1}$ , the  $L_{2}$  Wasserstein distance can be defined as
$$
W _ {2} \left(\mu_ {0}, \mu_ {1}\right) ^ {2} = \min  _ {\rho , v} \left\{\int \int \frac {1}{2} \rho (x, t) | v (x, t) | ^ {2} d x d t \right\} \tag {11}
$$
where the density  $\rho$  and the velocity  $v$  satisfy:
$$
\frac {d \rho (x , t)}{d t} = - \nabla \cdot (v (x, t) \rho (x, t)), v (x, t) = \nabla u (x, t) \tag {12}
$$
The optimality condition of the velocity is given by the generalized Hamilton-Jacobi (HJ) equation (i.e.,  $\partial_t u + \frac{1}{2} ||\nabla u||^2\leq 0$ ). The detailed derivation is deferred to the supplementary. We thus encourage our potential to satisfy the HJ equation with an external driving force as
$$
\frac {\partial}{\partial t} u ^ {k} (z, t) + \frac {1}{2} | | \nabla_ {z} u ^ {k} (z, t) | | ^ {2} = f (z, t) \quad \text {s u b j e c t t o} \quad f (z, t) \leq 0 \tag {13}
$$
Here we use another MLP to parameterize the external force  $f(z, t)$  and realize the negativity constraint by setting  $f(z, t) = -\mathrm{MLP}([z; t])^2$ . Notice that here we take the external force as learnable MLPs simply because we would like to obtain a flexible negativity constraint. The MLP architecture is set the same for both  $u(z, t)$  and  $f(z, t)$ . To achieve the PDE constraint, we impose a Physics-Informed Neural Network (PINN) [67] loss as
$$
\mathcal {L} _ {H J} = \frac {1}{T} \sum_ {t = 1} ^ {T} \left(\frac {\partial}{\partial t} u ^ {k} (\boldsymbol {z}, t) + \frac {1}{2} \left\| \nabla_ {\boldsymbol {z}} u ^ {k} (\boldsymbol {z}, t) \right\| ^ {2} - f (\boldsymbol {z}, t)\right) ^ {2} + \left\| \nabla u ^ {k} \left(\boldsymbol {z} _ {0}, 0\right) \right\| ^ {2} \tag {14}
$$
where the first term restricts the potential to obey the HJ equation, and the second term limits  $u(\mathbf{z}_t,t)$  to return no update at  $t = 0$ , therefore matching the initial condition.
This section starts with the experimental setup, followed by the main qualitative and quantitative results, then proceeds to discussions about the generalization ability to different composability and unseen data, and ends with the results on complex real-world datasets.
Figure 3: Exemplary latent evolution results of Scaling, Rotation, and Coloring on MNIST [54]. The top two rows are based on the supervised experiment, while the images of the bottom row are taken from the weakly-supervised setting of our experiment.
Datasets. We evaluate our method on two widely-used datasets in generative modeling, namely MNIST [54] and Shapes3D [10]. For MNIST [54], we manually construct three simple transformations including Scaling, Rotation, and Coloring. For Shapes3D [10], we use the self-contained four transformations that consist of Floor Hue, Wall Hue, Object Hue, and Scale.
Besides these two common benchmarks, we take a step further to apply our method on Falcol3D and Isaac3D [61], two complex large-scale and real-world datasets that contain sequences of different transformations. Falcol3D consists of indoor 3D scenes in different lighting conditions and viewpoints, while Isaac3D is a dataset of various robot arm movements in dynamic environments.
Figure 4: Exemplary latent flow results on Shapes3D [10]. The transformations from top to bottom are Floor Hue, Wall Hue, Object Hue, and Scale, respectively. The images of the top row are from the supervised experiment, while the bottom row is based on the weakly-supervised experiment.
Baselines. We mainly compare our method with SlowVAE [49] and Topographic VAE (TVAE) [45]. These two baselines could both achieve approximate equivariance. Specifically, TVAE introduces some learned latent operators, while SlowVAE enforces the Laplacian prior  $p(\pmb{z}_t|\pmb{z}_{t-1}) = \prod_{i=1}^{\alpha}\frac{\lambda}{2\Gamma(1/\alpha)}\exp(-\lambda|z_{t,i} - z_{t-1,i}|^\alpha)$  to sequential pairs. Within the disentanglement literature, our method is compared with the supervised PoFlow [79] which adopts a wave-like potential flow for sample evolution, and the unsupervised  $\beta$ -VAE [35] and FactorVAE [46] which encourage independence between single latent dimensions. Finally, the vanilla VAE is used as a controlled baseline.
Methods	Supervision?	Equivariance Error (↓)	Log-likelihood (↑)
Scaling	Rotation	Coloring
VAE [47]	No (✗)	1275.31±1.89	1310.72±2.19	1368.92±2.33	-2206.17±1.83
β-VAE [35]	No (✗)	741.58±4.57	751.32±5.22	808.16±5.03	-2224.67±2.35
FactorVAE [46]	No (✗)	659.71±4.89	632.44±5.76	662.18±5.26	-2209.33±2.47
SlowVAE [49]	Weak (✓)	461.59±5.37	447.46±5.46	398.12±4.83	-2197.68±2.39
TVAE [45]	Yes (✓)	505.19±2.77	493.28±3.37	451.25±2.76	-2181.13±1.87
PoFlow [79]	Yes (✓)	234.78±2.91	231.42±2.98	240.57±2.58	-2145.03±2.01
Ours	Yes (✓)	185.42±2.35	153.54±3.10	158.57±2.95	-2112.45±1.57
Ours	Weak (✓)	193.84±2.47	157.16±3.24	165.19±2.78	-2119.94±1.76
Table 1: Equivariance error  $\mathcal{E}_k$  and log-likelihood  $\log p(\pmb{x}_t)$  on MNIST [54].