# Formatting Instructions For NeurIPS 2024

Zhangkai Wu, ♡,♣
<EMAIL_ADDRESS>
&Xuhui Fan♣
<EMAIL_ADDRESS>
Jin Li♡
<EMAIL_ADDRESS>
&Zhilin Zhao♣
<EMAIL_ADDRESS>
&Hui Chen♣
<EMAIL_ADDRESS>
&Longbing Cao♣
<EMAIL_ADDRESS>
♡University of Technology Sydney
♣Macquire University

# ParamReL: Learning Parameter Space Representation via Progressively Encoding
Bayesian Flow Networks

Zhangkai Wu, ♡,♣
<EMAIL_ADDRESS>
&Xuhui Fan♣
<EMAIL_ADDRESS>
Jin Li♡
<EMAIL_ADDRESS>
&Zhilin Zhao♣
<EMAIL_ADDRESS>
&Hui Chen♣
<EMAIL_ADDRESS>
&Longbing Cao♣
<EMAIL_ADDRESS>
♡University of Technology Sydney
♣Macquire University

###### Abstract

The recently proposed Bayesian Flow Networks (BFNs) show great potential in
modeling parameter spaces, offering a unified strategy for handling
continuous, discretized, and discrete data. However, BFNs cannot learn high-
level semantic representation from the parameter space since common encoders,
which encode data into one static representation, cannot capture semantic
changes in parameters. This motivates a new direction: learning semantic
representations hidden in the parameter spaces to characterize mixed-typed
noisy data. Accordingly, we propose a representation learning framework named
ParamReL, which operates in the parameter space to obtain parameter-wise
latent semantics that exhibit progressive structures. Specifically, ParamReL
proposes a _self-_ encoder to learn latent semantics directly from parameters,
rather than from observations. The encoder is then integrated into BFNs,
enabling representation learning with various formats of observations. Mutual
information terms further promote the disentanglement of latent semantics and
capture meaningful semantics simultaneously. We illustrate conditional
generation and reconstruction in ParamReL via expanding BFNs, and extensive
quantitative experimental results demonstrate the superior effectiveness of
ParamReL in learning parameter representation.

##  1 Introduction

Representation learning [4], which aims at learning low-dimensional latent
semantics from high-dimensional observations, offers an unsupervised approach
to discovering high-level semantics in observations. It has been widely
applied in areas such as computer vision [26, 53, 10], and data analytics [41,
30]. While most representation learning methods [21, 6, 29] work on
continuous-valued observations, different non-trivial methods are needed to
discover semantics for the discretized [44, 32] and discrete data [2, 7].
Consequently, these individual efforts might face issues such as inconsistent
discoveries within the data [57] or repeated modelling efforts [23, 54].

On the other hand, Bayesian Flow Networks (BFNs) [16, 38, 50] have been
recently proposed as a promising deep generative model. By operating in the
parameter space, BFNs design a multi-step mechanism to approximate the ground-
truth parameters of observation sequentially. As a result, a uniform strategy
may be adopted to deal with continuous, discretized, and discrete data while
simultaneously maintaining fast sampling. Pilot studies of BFNs have shown
promising results in modelling different data formats.

Leveraging BFNs, this paper introduces ParamReL, a novel parameter space
representation learning framework that employs a unified strategy to extract
meaningful high-level semantics from continuous, discretized, and discrete
data. Specifically, a _self-_ encoder is designed to encode step-wise
parameters into lower-dimensional latent semantics, capturing gradual semantic
changes throughout the multi-step generation process. These latent semantics
are then integrated into a neural network architecture to form the parameters
for an output distribution that simulates observations. Furthermore, mutual
information is introduced to enhance the disentanglement of latent semantics,
promoting the capture of distinct and meaningful representations.

ParamReL is applied on benchmark datasets and verifies its effectiveness in
obtaining meaningful high-level semantics for discrete and continuous-valued
observations. Sampling and reverse-sampling procedures are developed here to
complete conditional image reconstruction and generation tasks. In particular,
our developed self-encoder discovers interesting progressive semantics along
with the flow steps. That is, our ParamReL obtains meaningful and clearer
disentangled representations while maintaining high sample quality.

The main contributions of this work can be summarized as follows: (1) A
parameter space representation learning framework introduces a uniform
strategy for modelling continuous, discretized and discrete observations; (2)
A _self-_ encoder encodes step-wise parameters into step-wise semantics to
reveal a series of gradually changing latent semantics; (3) A mutual
information term promotes latent semantics being disentangled and storing
meaningful semantics simultaneously; (4) Sampling and reverse-sampling methods
are developed, and generation and reconstruction tasks are completed in the
parameter space.

##  2 Bayesian Flow Networks

Bayesian Flow Networks (BFNs) [16, 38, 50] are a new class of latent variable
models. Similar to the reverse process of diffusion models [18, 37], BFNs
target optimizing the neural network architecture on posterior parameters to
simulate observations as they are gradually observed. In particular, BFNs
assume two types of distributions: a simple _input distribution_
$P_{\mathrm{I}}(\cdot)$ representing the initial belief about observations and
an _output distribution_ $P_{\mathrm{O}}(\cdot)$ representing the simulating
distribution for observations. The parameters of input distribution are first
updated through a Bayesian inference scheme and then passed into a neural
network to form the parameters of output distributions. The main objective of
BFNs is to minimize the Kullback-Leibler (KL) divergence between the ground-
truth data distribution and the output distribution, ensuring that the output
distribution closely approximates the ground-truth data distribution.

Following the notations in diffusion models, we use
${\boldsymbol{\mathbf{x}}}_{0}$ to denote the observations. There are $T$
_reverse_ steps in BFNs which gradually reveals the information of
${\boldsymbol{\mathbf{x}}}_{0}$ through
$\\{{\boldsymbol{\mathbf{x}}}_{T},{\boldsymbol{\mathbf{x}}}_{T-1},\ldots,{\boldsymbol{\mathbf{x}}}_{1}\\}$
to the input distribution111It is noted that the index $t$ is used reversely
in [16]. We make such changes to be consistent with the diffusion model
settings [18, 37].. At the $t$-th step, the reverse step is denoted as a
_sender distribution_
$p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})$,
where $\alpha_{t}$ is a precision parameter. Combined with input distribution
$p_{\mathrm{I}}({\boldsymbol{\mathbf{x}}}_{t};{\boldsymbol{\mathbf{\theta}}}_{t+1})$,
the posterior distribution of ${\boldsymbol{\mathbf{x}}}_{t}$ is obtained as
$p({\boldsymbol{\mathbf{x}}}_{t};{\boldsymbol{\mathbf{\theta}}}_{t}=h({\boldsymbol{\mathbf{\theta}}}_{t+1},{\boldsymbol{\mathbf{x}}}_{t-1},\alpha_{t}))\propto
p_{\mathrm{I}}({\boldsymbol{\mathbf{x}}}_{t};{\boldsymbol{\mathbf{\theta}}}_{t+1})p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})$,
where
$h({\boldsymbol{\mathbf{\theta}}}_{t+1},{\boldsymbol{\mathbf{x}}}_{t-1},\alpha_{t})$
is the Bayesian update function. By feeding this posterior parameter
${\boldsymbol{\mathbf{\theta}}}_{t}$ into a neural network $\psi(\cdot)$,
${\boldsymbol{\mathbf{x}}}_{t}$’s output distribution $p_{\mathrm{O}}(\cdot)$
is parameterized as
$p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t}))$.
Finally, a _receiver distribution_ $p_{\mathrm{R}}(\cdot)$ is defined as the
expectation of the sender distribution with respect to the output
distribution, i.e.,
$p_{\mathrm{R}}({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t}),\alpha_{t}):=\mathbb{E}_{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t}))}[p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})]$.
See Figure 1 (a) for a visualization of the relationships between these
distributions.

In BFNs, the joint distribution over the observation
${\boldsymbol{\mathbf{x}}}_{0}$ and the intermediates
$\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}$ is defined as
$p({\boldsymbol{\mathbf{x}}}_{0},\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}|-):=p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0}))\prod_{t=1}^{T}p_{\mathrm{R}}({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t}),\alpha_{t})$.
This intractable joint distribution can be approximated under the variational
inference framework as follows:

$\log
p({\boldsymbol{\mathbf{x}}}_{0})\geq\mathbb{E}_{p_{\mathrm{F}}({\boldsymbol{\mathbf{\theta}}}_{1:T}|-)p_{\mathrm{S}}(\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}|-)}\left[\log\frac{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0}))\prod_{t=1}^{T}p_{\mathrm{R}}({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t}),\alpha_{t})}{\prod_{t=1}^{T}p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})}\right]\\\
=-\sum_{t=1}^{T}\underbrace{\mathbb{E}_{p_{F}({\boldsymbol{\mathbf{\theta}}}_{t}|-)}\textsc{kl}\left[p_{\mathrm{S}}\left({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{{\boldsymbol{\mathbf{x}}}_{0}};\alpha_{T:t}\right)\;\|\;p_{\mathrm{R}}\left({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t}),\alpha_{t}\right)\right]}_{\mathcal{L}_{t}^{\mathrm{R}}({\boldsymbol{\mathbf{x}}})}+\underbrace{\mathbb{E}_{p_{F}({\boldsymbol{\mathbf{\theta}}}_{0}|-)}\ln
p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0}))}_{\mathcal{L}^{{\boldsymbol{\mathbf{D}}}}({\boldsymbol{\mathbf{x}}})},$
(1)

where $p_{\mathrm{F}}({\boldsymbol{\mathbf{\theta}}}_{t}|-)$ represents the
distribution of ${\boldsymbol{\mathbf{\theta}}}_{t}$ (see Appendix A.1 for a
detailed calculation). Maximizing Eq. 1 equals minimizing the discrepancy
between the sender and receiver distributions and implicitly penalizing
Distortion
$\mathcal{L}^{{\boldsymbol{\mathbf{D}}}}({\boldsymbol{\mathbf{x}}})$ to
maximize the likelihood distribution over data.

By working in the parameter space, BFNs can uniformly model continuous,
discretized, and discrete observations. For example, BFNs can use the mean
parameters of Gaussian distributions to model continuous data and the event
probability parameters of categorical distributions to study discrete data
(see detailed settings for distributions in Table 1). However, BFNs cannot
produce meaningful latent semantics that discover high-level concepts in
observations, such as hair colors in portrait images.

Table 1: Examples of detailed distribution formats in BFNs.

Data type | $p_{\mathrm{I}}({\boldsymbol{\mathbf{x}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t+1})$ | $p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}|{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})$ | ${\boldsymbol{\mathbf{\theta}}}_{t}=h({\boldsymbol{\mathbf{\theta}}}_{t+1},{\boldsymbol{\mathbf{x}}}_{t-1},\alpha_{t})$
---|---|---|---
Continuous data | $\mathcal{N}({\boldsymbol{\mathbf{x}}}_{t};{\boldsymbol{\mathbf{\theta}}}_{t+1}=\\{\mu_{t+1},\rho_{t+1}^{-1}\\})$ | $\mathcal{N}({\boldsymbol{\mathbf{x}}}_{t-1};{\boldsymbol{\mathbf{x}}},\alpha_{t}^{-1})$ | $\mu_{t}=\frac{\alpha_{t}{\boldsymbol{\mathbf{x}}}_{t-1}+\rho_{t+1}\mu_{t+1}}{\alpha_{t}+\rho_{t+1}}$
Discrete data | $\text{categorical}(\text{x};\frac{{\boldsymbol{\mathbf{1}}}}{K})$ | $\mathcal{N}({\boldsymbol{\mathbf{x}}}_{t-1};\alpha_{t}(K{\boldsymbol{\mathbf{e}}}_{{\boldsymbol{\mathbf{x}}}_{t}}-1),\alpha_{t}K{\boldsymbol{\mathbf{I}}})$ | ${\boldsymbol{\mathbf{\theta}}}_{t}=\frac{e^{{\boldsymbol{\mathbf{x}}}_{t-1}}{\boldsymbol{\mathbf{\theta}}}_{t+1}}{\sum_{k}e^{{\boldsymbol{\mathbf{x}}}_{t-1,k}}\theta_{t+1,k}}$
Data type | $p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t})$ | $p_{\mathrm{R}}({\boldsymbol{\mathbf{x}}}_{t-1}|\psi({\boldsymbol{\mathbf{\theta}}}_{t}),\alpha_{t})$
Continuous data | $\delta({\boldsymbol{\mathbf{x}}}_{t}-\psi({\boldsymbol{\mathbf{\theta}}}_{t}))$ | $\mathcal{N}({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t}),\alpha_{t}^{-1})$
Discrete data | $\text{categorical}(\text{softmax}(\psi({\boldsymbol{\mathbf{\theta}}}_{t})))$ | $\sum_{k}p_{O}(k;\psi({\boldsymbol{\mathbf{\theta}}}_{t}))\mathcal{N}({\boldsymbol{\mathbf{x}}}_{t-1};\alpha_{t}(K{\boldsymbol{\mathbf{e}}}_{k}-1),\alpha_{t}K{\boldsymbol{\mathbf{I}}})$

Figure 1: The relationships between distributions in BFNs (a) and ParamReL
(b).

##  3 ParamReL: Parameter Space Representation Learning Framework

This section introduces the ParamReL framework, which facilitates parameter
space representation learning through low-dimensional latent semantics.
Instead of approximating the data distribution
$p({\boldsymbol{\mathbf{x}}}_{0})$, ParamReL aims to learn the joint
distribution over data ${\boldsymbol{\mathbf{x}}}_{0}$ and latent semantics
${\boldsymbol{\mathbf{z}}}$, with
$|{\boldsymbol{\mathbf{z}}}|\ll|{\boldsymbol{\mathbf{x}}}_{0}|$. This ensures
that low-dimensional latent semantics ${\boldsymbol{\mathbf{z}}}$ are learned
from high-dimensional data ${\boldsymbol{\mathbf{x}}}_{0}$. In other words,
ParamReL seeks to reconstruct observations ${\boldsymbol{\mathbf{x}}}_{0}$
using a high-performing likelihood function while obtaining meaningful low-
dimensional latent semantics ${\boldsymbol{\mathbf{z}}}$.

Figure 2 shows the framework of ParamReL, which involves $T$ steps of Bayesian
updates, parameter encoding, and latent decoding. At step
$t=\\{T,T-1,\ldots,1\\}$, the observations first go through a Bayesian update
function and a sender distribution to form the posterior

Figure 2: The framework of ParamReL, consisting of a semantic encoder
$q_{\phi}$ and a parameter decoder $p_{\Psi}$. During the reverse stage,
parameters of data distribution $\theta_{t}$ will encode a time-specific
semantic latent ${\boldsymbol{\mathbf{z}}}_{t}$ and then decode the $t-1$-th
parameters $\theta_{t-1}$.

###  3.1 Self-encoder
$q_{\mathbf{\phi}}(\mathbf{z}_{t}|\mathbf{\theta}_{t},t)$

Common autoencoders, such as
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}|{\boldsymbol{\mathbf{x}}}_{0})$,
aim to encode the observation ${\boldsymbol{\mathbf{x}}}_{0}$ into a static
latent semantic ${\boldsymbol{\mathbf{z}}}$. However, this construction does
not align well with the ParamReL framework, which employs multiple steps to
generate observations. In ParamReL, a single static semantic
${\boldsymbol{\mathbf{z}}}$ may not be sufficient to encode all intermediate
variables $\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}$. Specifically,
${\boldsymbol{\mathbf{z}}}$ works together with the updated parameter
${\boldsymbol{\mathbf{\theta}}}_{t}$ in
$\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}})$. A static
semantic ${\boldsymbol{\mathbf{z}}}$ derived from the initial observation
${\boldsymbol{\mathbf{x}}}_{0}$ may fail to capture the complex dynamics in
$\\{{\boldsymbol{\mathbf{\theta}}}_{t}\\}_{t}$.

In ParamReL, we propose a _self-encoder_
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$
that differs from the standard encoder
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}|{\boldsymbol{\mathbf{x}}}_{0})$
in two key ways: (1) We adopt a step-wise representation
${\boldsymbol{\mathbf{z}}}_{t}$ tailored to the variable behaviours at each
step $t$; (2) We utilize ${\boldsymbol{\mathbf{\theta}}}_{t}$, which
summarizes information from all previous steps, instead of
${\boldsymbol{\mathbf{x}}}_{t}$, which only captures the current step $t$, to
generate the step-wise ${\boldsymbol{\mathbf{z}}}_{t}$. Through
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$,
${\boldsymbol{\mathbf{\theta}}}_{t}$ effectively encodes _itself_ into
${\boldsymbol{\mathbf{z}}}_{t}$, and they then jointly form
$\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t})$.
Therefore, we refer to it as a _self-encoder_.

The proposed self-encoder works consistently with the decoder network here as
both work on ${\boldsymbol{\mathbf{\theta}}}_{t}$, see Figure 1 (b). By
obtaining _a series of_ latent semantics
$\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}$, it is expected these encoded
latents would show progressive semantic behaviour (representing, for example,
the gradual changes of age _and_ smiling _and_ skin colour of a person) along
with the generation procedures. This is followed by a reconstruction phase,
where lower-level generative latent ${{\boldsymbol{\mathbf{x}}}}_{t}$ (such as
hair texture) are progressively incorporated. In our experiments, we
parameterize
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$
as a UNet encoder (see Appendix C for more details).

Prior for ${\boldsymbol{\mathbf{z}}}_{t}$: Ideally, the latent semantics
$\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}$ should provide low-dimensional
representations that are different from the intermediates
$\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}$ of BFNs, yet they should not
compromise the data reconstruction process. To generate high-quality semantic
representations, a learnable and smooth latent space is essential [4]. This
necessitates the integration of the prior distribution
$p({\boldsymbol{\mathbf{z}}}_{t})$ into a robust probabilistic framework,
thereby enabling the unconditional sampling of
${{\boldsymbol{\mathbf{x}}}}_{0}$. For simplicity and sampling efficiency,
$p({\boldsymbol{\mathbf{z}}}_{t})$ could be assumed to follow a
straightforward distribution, such as a Gaussian.

###  3.2 Training and Testing with ParamReL

By augmenting the latents $\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}$, the joint
distribution of
${\boldsymbol{\mathbf{x}}}_{0},\\{{{\boldsymbol{\mathbf{x}}}}_{t}\\}_{t},\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}$
can be written as:

$\displaystyle
p\left({{\boldsymbol{\mathbf{x}}}}_{0},\\{{{\boldsymbol{\mathbf{x}}}}_{t}\\}_{t},\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}\,|\,{\boldsymbol{\mathbf{\theta}}}_{0},\alpha\right)=p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))\prod_{t=1}^{T}p({\boldsymbol{\mathbf{z}}}_{t})\mathbb{E}_{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}))}[p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})],$
(2)

where the model implements an output distribution
$p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}))$
conditioned on auxiliary latent ${{\boldsymbol{\mathbf{z}}}}_{t}$, which is
distributed according to prior $p({{\boldsymbol{\mathbf{z}}}}_{t})$. The
latent variable ${{\boldsymbol{\mathbf{z}}}}_{t}$ of prior
$p({{\boldsymbol{\mathbf{z}}}}_{t})$ is independent of the BFNs process as a
latent representation of input. Through a neural network architecture
$\psi(\cdot,\cdot)$, ${\boldsymbol{\mathbf{z}}}_{t}$ works together with the
posterior parameter ${\boldsymbol{\mathbf{\theta}}}_{t}$ to form the
parameters of the output distribution as
$\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t})$. We
use
$p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}))\prod_{t=1}^{T}p_{\mathrm{R}}({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}),\alpha_{t})$
to decode both the auxiliary latent representation and the generative latent.

Training with ParamRL: Variational inference methods approximate the
intractable joint distribution. With
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}\,|\,{\boldsymbol{\mathbf{\theta}}}_{t},t)$
defined as the encoder for ${\boldsymbol{\mathbf{z}}}_{t}$ and
$p_{\mathrm{S}}(\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}|-)$ defined as the
variational distribution for ${\boldsymbol{\mathbf{x}}}_{T:1}$, the evidence
lower bound on the marginal log-likelihood of data
${\boldsymbol{\mathbf{x}}}_{0}$ is (see the full derivation in Appendix B):

$\log
p({\boldsymbol{\mathbf{x}}}_{0})\geq-\sum_{t=1}^{T}\mathbb{E}_{p_{\mathrm{F}}({\boldsymbol{\mathbf{\theta}}}_{t}|-)}\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t})}\left\\{{\textsc{kl}\left[p_{\mathrm{S}}\left({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{{\boldsymbol{\mathbf{x}}}_{0}};\alpha_{T:t}\right)\;\|\;p_{\mathrm{R}}\left({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}),\alpha_{t}\right)\right]}\right.\\\
\left.-\textsc{kl}\left[q_{{\boldsymbol{\mathbf{\phi}}}}\left({{\boldsymbol{\mathbf{z}}}}_{t}\,|\,{{\boldsymbol{\mathbf{\theta}}}_{t}},t\right)\;\|\;p({{\boldsymbol{\mathbf{z}}}}_{t})\right]\right\\}+\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{0},{\boldsymbol{\mathbf{\theta}}}_{0})}\left[{\ln
p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))}\right]:=\mathcal{L}_{\texttt{ParamReL}}.$
(3)

Maximizing the objective function $\mathcal{L}_{\text{ParamReL}}$ is
equivalent to performing amortized inference [22] through encoders
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$
and learning likelihood function through decoders [55]. The encodable
posterior
$q_{{\boldsymbol{\mathbf{\phi}}}}\left({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t\right)$
is used to infer high-level semantics
$\\{{{\boldsymbol{\mathbf{x}}}}_{t}\\}_{t}$. During the _decoding_ stage,
${{\boldsymbol{\mathbf{x}}}}_{0}$ is used into generative latent variables
${\boldsymbol{\mathbf{x}}}_{1:T}$, which contains low-level semantic features
in generating the observations. In ParamReL, the parameters of the output
distribution are learned through iteratively proceeding the Bayesian updating
process and a learned noise model
${\psi}({{\boldsymbol{\mathbf{\theta}}}},{\boldsymbol{\mathbf{z}}})$
parameterized by neural networks $\psi$.

Sampling and Reverse-sampling with ParamReL in the Test Phase Once ParamReL is
trained, we can use the generative procedure to do pseudo-data sampling, i.e.,
use ${\boldsymbol{\mathbf{x}}}_{T}$ to sequentially generate
${\boldsymbol{\mathbf{x}}}_{T-1},\dots,{\boldsymbol{\mathbf{x}}}_{0}$ from
${\boldsymbol{\mathbf{x}}}_{T}$. In particular, at the $t$-th step, we have

$\displaystyle\widehat{{\boldsymbol{\mathbf{x}}}}_{t-1}\sim
p_{\mathrm{S}}\left(\widehat{{\boldsymbol{\mathbf{x}}}}_{t-1}|{\boldsymbol{\mathbf{x}}}_{t},\alpha_{t}\right),\
{\boldsymbol{\mathbf{\theta}}}_{t}=h({\boldsymbol{\mathbf{\theta}}}_{t+1},\widehat{{\boldsymbol{\mathbf{x}}}}_{t-1},\alpha_{t}),\
{\boldsymbol{\mathbf{z}}}_{t}\sim p({\boldsymbol{\mathbf{z}}}_{t}),\
\widehat{{\boldsymbol{\mathbf{x}}}}_{t}\sim
p_{O}(\widehat{{\boldsymbol{\mathbf{x}}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t})),\
{\boldsymbol{\mathbf{x}}}_{t-1}=\widehat{{\boldsymbol{\mathbf{x}}}}_{t}$ (4)

where $\widehat{{\boldsymbol{\mathbf{x}}}}_{t-1}$ is a pseudo value for the
intermediate ${\boldsymbol{\mathbf{x}}}_{t-1}$. Since the self-encoder is
already trained, it may be possible to use
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$
to replace ${\boldsymbol{\mathbf{z}}}_{t}$’s prior
$p({\boldsymbol{\mathbf{z}}}_{t})$ to improve the sampling quality.

However, the reverse-sampling procedure, which transits observation
${\boldsymbol{\mathbf{x}}}_{0}$ through
${\boldsymbol{\mathbf{x}}}_{1},{\boldsymbol{\mathbf{x}}}_{2},\ldots$ until its
noisy state ${\boldsymbol{\mathbf{x}}}_{T}$, is not as straightforward as the
sampling procedure. In fact, we notice the output distribution
$p_{O}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}))$
may help us to derive the reverse-sampling procedure. At the $t$-th step, the
reverse-sampling procedure may be implemented as

$\displaystyle{\boldsymbol{\mathbf{\theta}}}_{t+1}=h^{-1}({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{x}}}_{t-1},\alpha_{t}),\quad{\boldsymbol{\mathbf{z}}}_{t+1}=q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{\theta}}}_{t+1},t+1),\quad{\boldsymbol{\mathbf{x}}}_{t+1}\sim
p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t+1};\psi({\boldsymbol{\mathbf{\theta}}}_{t+1},{\boldsymbol{\mathbf{z}}}_{t+1})).$
(5)

where $h^{-1}(\cdot)$ refers to the inverse Bayesian update function. Given
the sampling and reverse-sampling procedure, we can do various tasks such as
image reconstruction, interpolation, etc. The details of such results are
shown in the experiments.

###  3.3 Regularizing Semantics By Maximizing Mutual Information

Ideally, during the training phase, we want to acquire the semantic
representation ${\boldsymbol{\mathbf{z}}}_{t}$ by the self-encoder
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$
and achieve high-quality reconstruction
$\widehat{{\boldsymbol{\mathbf{x}}}_{0}}$ by the decoder (i.e., the output
distribution
$p_{O}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))$).
However, there exists a trade-off between inference and learning [34, 46]
coherent in optimizing the objective function $\mathcal{L}_{\text{ParamReL}}$
in Eq. (3). In most cases, ParamReL favours fitting likelihood rather than
inference [56]. Based on the rate-distortion theory [1, 3], the rate (KL
divergence term) constrained by the encoders compresses enough information for
the distortion (reconstruction error) but less informative to induce a smooth
latent space.

To remedy the insufficient representation learning during the inference stage,
we want to increase the dependence between input parameters
${\boldsymbol{\mathbf{\theta}}}_{t}$ and encoded latent
${\boldsymbol{\mathbf{z}}}_{t}$ by maximizing the mutual information
$I({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t})$.

We can rewrite the tractable learning object in ParamReL by adding the mutual
information maximization term as:
$\mathcal{L}_{\text{ParamReL+}}=\mathcal{L}_{\text{ParamReL}}+\frac{\gamma}{T}\sum_{t}I_{q}({\boldsymbol{\mathbf{\theta}}}_{t};{\boldsymbol{\mathbf{z}}}_{t})$,
where
$I_{q}({\boldsymbol{\mathbf{\theta}}}_{t};{\boldsymbol{\mathbf{z}}}_{t})$ is
the mutual information between
${\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}$ under the
distribution
$q_{\phi}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t})$.
Considering that we cannot optimize this object directly, we can rewrite this
object by factorizing the rate term into MI and Total Correlation (TC):

$\mathcal{L}_{\texttt{ParamReL+}}=-\sum_{t=1}^{T}\mathbb{E}_{p_{\mathrm{F}}({\boldsymbol{\mathbf{\theta}}}_{t}|-)}\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t})}\left\\{{\textsc{kl}\left[p_{\mathrm{S}}\left({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{{\boldsymbol{\mathbf{x}}}_{0}};\alpha_{T:t}\right)\;\|\;p_{\mathrm{R}}\left({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}),\alpha_{t}\right)\right]}\right.\\\
\left.-\frac{1-\gamma}{T}\textsc{kl}\left[q_{{\boldsymbol{\mathbf{\phi}}}}\left({{\boldsymbol{\mathbf{z}}}}_{t}\,|\,{{\boldsymbol{\mathbf{\theta}}}_{t}}\right)\;\|\;p({{\boldsymbol{\mathbf{z}}}})\right]-\frac{\gamma+\lambda-1}{T}\textsc{kl}\left[q_{{\boldsymbol{\mathbf{\phi}}}}\left({{\boldsymbol{\mathbf{z}}}}_{t}\right)\;\|\;p({{\boldsymbol{\mathbf{z}}}})\right]\right\\}+\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{0},{\boldsymbol{\mathbf{\theta}}}_{0})}\left[{\ln
p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))}\right].$
(6)

Mutual Information Learning Unlike the rest of the terms that can be optimized
directly using reparameterization tricks, the TC term cannot be directly
optimized due to intractable marginal distribution
$q_{{\boldsymbol{\mathbf{\phi}}}}({{\boldsymbol{\mathbf{z}}}_{t}})$. Here, we
follow the guidance in [56] to replace the TC term with any strict divergence
$D$, where
$D\left(q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}})\|p({\boldsymbol{\mathbf{z}}})\right)=0$
iff
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}})=p({\boldsymbol{\mathbf{z}}})$.
We implement the Maximum-Mean Discrepancy (MMD) [56] from the divergence
family. MMD is a statistical measure that quantifies the difference between
two probability distributions by comparing their mean embeddings in a high-
dimensional feature space. By defining the kernel function
$\kappa(\cdot,\cdot)$, $D_{\mathrm{MMD}}$ is denoted as:

$D_{\mathrm{MMD}}\left(q(\cdot)\|p(\cdot)\right)=\mathbb{E}_{p({\boldsymbol{\mathbf{z}}}),p\left({\boldsymbol{\mathbf{z}}}^{\prime}\right)}\left[\kappa\left({\boldsymbol{\mathbf{z}}},{\boldsymbol{\mathbf{z}}}^{\prime}\right)\right]-2\mathbb{E}_{q({\boldsymbol{\mathbf{z}}}),p\left({\boldsymbol{\mathbf{z}}}^{\prime}\right)}\left[\kappa\left({\boldsymbol{\mathbf{z}}},{\boldsymbol{\mathbf{z}}}^{\prime}\right)\right]+\mathbb{E}_{q({\boldsymbol{\mathbf{z}}}),q\left({\boldsymbol{\mathbf{z}}}^{\prime}\right)}\left[\kappa\left({\boldsymbol{\mathbf{z}}},{\boldsymbol{\mathbf{z}}}^{\prime}\right)\right].$
(7)

##  4 Related Work

Recent advances have demonstrated that diffusion-based models [18, 37] are
capable of generating high-quality data. Nonetheless, compared to the
autoencoder framework, the intermediate outputs in diffusion stages are high-
dimensional and lack smoothness, making them unsuitable for representation
learning. Contemporary research focuses on encoding a conditional latent space
to acquire semantic low-dimensional representations. However, sample-based
models [31, 45], such as VAEs and Diffusions, exhibit limitations when applied
to discrete data.

Deep hierarchical VAEs have seen progress in capturing latent dependence
structures for encoding an expressive posterior, statistically or
semantically. VQVAE-based [44, 32] models have local-to-global feature-based
explanatory hierarchies at the image level, forming a codebook-based discrete
posterior. [36, 40] build recursive latent structures in multi-layer networks
to form an aggregated posterior. NVAE [43] demonstrates that depth-wise
hierarchies encoded by residual networks can approximate the posterior
precisely despite using shallow networks. Unlike the sample-based encoder,
where the information flow between input and latent is maximized in encoding-
decoding pipelines on the sample space, we develop progressive encoders in the
parameter space to capture the dynamic semantics.

##  5 Experiments

We present two ParamReL variants in different parameter spaces: for the
continuous input distribution, denoted as ParamReLc, and for the discrete
input distribution, denoted as ParamReLd; refer to Table 1 for the
distribution definitions. We evaluate ParamReLs on discrete and continuous
benchmarks, including three discrete datasets binarized by grey-style data:
binarized MNIST (bMNIST) [8], binarized FashionMNIST (bFashionMNIST) [47]; as
well as three continuous datasets: CelebA [27], CIFAR10 [24], and Shapes3d [5]
222For the discrete version, continuous data ($k$-bit images) can be
discretized into $2^{k}$ bins. This process involves dividing the data range
$[-1,1]$ into $k$ intervals, each with a length of $2/k$..The hyperparameter
choices and experiment configurations corresponding to each dataset for
training can be found in Appendix D.

We evaluate the representation learning ability in three reconstruction-based
tasks, i.e., decoding the samples from the encoder and the decoder, including
latent interpolation, disentanglement, and time-varying conditional
reconstruction; and one generation-based task, i.e., decoding the samples only
from the decoder with a given prior, including conditional generation.

We conduct a comprehensive two-fold comparison. Firstly, we compare our
parameter-based variants against sample-based representation learning
baselines, including AE and VAE-based models: $\beta$-VAE [17], infoVAE [55],
and diffusion-based models: DiffusionAE [31], infoDiffusion [45]. These models
represent significant advancements in representation learning, with
$\beta$-VAE being the pioneer in VAE-based disentanglement, infoVAE
introducing MMD to strike a balance between generation and representation,
DiffusionAE integrating AEs into Diffusion models for encodable latent
learning, and infoDiffusion incorporating VAEs into diffusion models for
disentanglement learning. Secondly, we compare the ParamReLd and ParamReLc
variants across various input distributions for continuous and discrete data
representation learning. This comparison enables a detailed investigation of
parameter space assumptions for discrete and continuous representation
learning.

###  5.1 Representation Learning on Discrete Data by ParamReLd

Figure 3: The left panel (a-b) shows high-level semantic latent captured by
${\boldsymbol{\mathbf{z}}}_{\mathrm{sem}}$ from ParamReL’s encoders. By fixing
${\boldsymbol{\mathbf{z}}}_{\mathrm{sem}}$, the global characters of the
images are invariant. By varying the stochastic subcodes
${\boldsymbol{\mathbf{x}}}_{0}$, the local attributes in the corresponding
generated images may vary, such as the Narrow_Eyes attribute in (a), the
Blond_Hair attribute in (b), and the Mouth_Slightly_Open attribute in (c). The
right panel (d-f) illustrates the time-varying changes that ParamReL’s
progressive encodes interfaced. By varying time encodes at 100,200,300 time
steps, more attributes will be influenced in the reconstruction stage: the
Big_Lips, Pointy_Nose attributes in (d), the Blond_Hair, Bald attributes in
(e) and the Wavy_Hair, High_Cheekbones attributes in (f).

####  5.1.1 Semantic Latent for Downstream Tasks

The representation ${\boldsymbol{\mathbf{z}}}_{\mathrm{sem}}$ should be
general and transferable [14] for downstream tasks. We measure the quality of
learned representations in classification tasks. To ensure their universality,
we assess the quality of the learned representations on various datasets by
different classifiers. Specifically, Following [48], we train a classifier on
labelled test sets for each ParamReL model. We allocate $80\%$ of the dataset
for training a classifier and reserve the remaining $20\%$ for testing
purposes. The performance on the test set is evaluated based on AUROC. This
process is conducted in a 5-fold cross-validation manner, with the results
reported as mean metrics ± one standard deviation. The results are shown in
Table 2. Higher AUROC would suggest that the learned features
${\boldsymbol{\mathbf{z}}}_{\mathrm{sem}}$ contain more information about data
${\boldsymbol{\mathbf{x}}}$. In addition to assessing the representation
quality, we also compare the image reconstruction ability against baselines.
From Table 2, we can conclude that VAE-based models still produce blurry
reconstructions, while diffusion-based and parameter-based models can build
near-exact reconstructions. Refer to Figure 9 in Appendix G for reconstructed
binary images.

###  5.2 Representation Learning on Continuous Data by ParamReLs

####  5.2.1 High-level Representation Learning for Conditional Generation

In Figure 3, it is demonstrated that high-level semantic information is
captured by the auxiliary variables ${\boldsymbol{\mathbf{z}}}_{\mathrm{sem}}$
in image generation tasks. This is illustrated by a set of latent variables
$<\mathbf{z}_{\mathrm{sem}}^{i},\mathbf{x}_{\mathrm{0}}^{i,j}>$, where the
auxiliary variables are fixed and encoded from the $i$-th input image by
trained the ParamReL’s encoders. Additionally, stochastic subcodes
${\boldsymbol{\mathbf{x}}}_{0}$ are sampled $j$ times from $\mathcal{N}(0,1)$
corresponding to the $i$-th input image. Concurrently, the low-level semantic
features, such as local attributes in images (e.g., Wearing_Hat, Big_Nose),
are captured by stochastic subcodes in the BFN decoders of ParamReL.

####  5.2.2 Time-varying Representation Learning for Conditional
Reconstruction

The effectiveness of progressive latent learned by the _self-encoder_ will be
validated on a new time-varying reconstruction task. Given input
${\boldsymbol{\mathbf{x}}}^{i}$, the time encoding latent set will be acquired
by encoding the $t$-th time and parameters, i.e.,
$\\{{\boldsymbol{\mathbf{z}}}_{0},...,{\boldsymbol{\mathbf{z}}}_{1000}\\}\sim
q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$.
We will reconstruct samples with time-specific semantic encoding and sample-
specific subcodes, along with the fixed subcode reversed from the same input
${\boldsymbol{\mathbf{x}}}^{i}$. In that case, the attributes will vary due to
the semantics evolution encoded by time-specific latent. Refer to Figure 3 for
more explanation.

Table 2: Comparison of representation learning algorithms on discrete data by classification accuracy and generation performance. ($\uparrow$ higher is better; $\mathit{g}$ for the Gaussian distribution, $\mathit{c}$ for the categorical distribution; [Top-1,Top-2], S for sample-based models, P for parameter-based models). |  |  | bMNIST | bFASHIONMNIST
---|---|---|---|---
|  Prior |  |  AUROC$\uparrow$ |  NotBlur? |  AUROC$\uparrow$ |  NotBlur?
S |  - |  AE |  - |  ✖ |  0.819±0.003 |  ✖
$\mathit{g}$ |  VAE |  - |  ✖ |  0.796±0.002 |  ✖
$\mathit{g}$ |  $\beta$-VAE |  0.842±0.459 |  ✖ |  0.779±0.004 |  ✖
$\mathit{g}$ |  infoVAE |  0.847±0.386 |  ✖ |  0.807±0.003 |  ✖
$\mathit{g}$ |  DiffAE |  - |  ✔ |  0.835±0.002 |  ✔
$\mathit{g}$ |  infoDiffusion |  0.898±0.430 |  ✔ |  0.839±0.003 |  ✔
P |  $\mathit{c}$ |  ParamReLd($\lambda=0.1$) |  0.927±0.086 |  ✔ |  0.857±0.666 |  ✔
$\mathit{c}$ |  ParamReLd (w ProgEnc) |  0.946±0.334 |  ✔ |  0.892±0.787 |  ✔

####  5.2.3 Smooth Representation Learning for Latent Interpolation

Latent space interpolation [15, 17] can validate the smoothness, continuity,
and semantic coherence in the learned representation learning space of
generative models. Generally, we can embed two samples into the corresponding
latent space, and interpolating between latent variables sampled from the
corresponding space yields interpolated representations. Their reconstruction
learned by generative models reveals the semantic richness of the learned
space. The specific latent interpolation process is detailed in Appendix E.

In Figure 4, in contrast to VAE variants (a) vanilla VAE, and (b) $\beta$-VAE,
ParamReL allows near-exact reconstruction. Compared with diffusion models (c)
DiffusionAE and (d) infoDiffusion, ParamReL has a smoother and more consistent
latent space with high-quality samples.

Figure 4: Comparisons of latent space interpolation among sample-based models
and parameter-based models on dataset CelebA. Only our ParamRel model (e) can
learn a continuous, smooth latent space while ensuring near-exact image
reconstruction. Specifically, while sample-based generative models can learn a
continuous but unsmooth latent space, this leads to incomplete
reconstructions. For example, in (a-d), the attribute of eyeglasses is
frequently omitted. Moreover, VAEs (a,b) tend to produce blurry images.
Additionally, it is observable that sample-based models often compromise
reconstruction in favour of representation learning, as evidenced by the
failure of diffusion model variants (c-d) to accurately reconstruct background
characters in imageB.

####  5.2.4 Disentanglement

We conduct latent traversals on the CeleA dataset to assess the
disentanglement properties of our trained ParamReL, as depicted in Figure 5.
Specifically, we modify one dimension at a time, i.e., $d(0\leq d\leq L-1)$,
within the semantic latent vector
${\boldsymbol{\mathbf{z}}}_{\mathrm{sem}}\in\mathbb{R}^{L}$, and then varying
it $M$ times in a standardized range (e.g., -3 to +3) to acquire traversed
latent ${\boldsymbol{\mathbf{z}}}_{\mathrm{traversal}}\in\mathbb{R}^{M\times
L}$, while keeping other dimensions constant. Upon decoding these adjusted
traversed latent variables, the generated samples are evaluated for changes in
specific attributes. Successful disentanglement is demonstrated when the
manipulation of one latent dimension changes a single distinguishable
attribute in the dataset, such as age, with all other attributes remaining
unchanged. This method provides a straightforward validation of the model’s
disentanglement capabilities.

Figure 5 reveals that ParamReL can isolate and control data attributes in an
unsupervised manner. This independence in attribute variation underscores the
effectiveness of a model in achieving disentanglement. Specifically, the
reconstructed sample changes independently in some semantic attribute, e.g.,
Smiling in Figure 5 (a), Pale Skin in Figure 5 (b), Big Nose in Figure 5 (c)
without affecting others.

Figure 5: Traversals of latent by ParamReL on CelebA. The interpretable
traversal directions are displayed by traversing the encodings ranging from
$[-3,3]$.

To facilitate a thorough and unbiased quantitative assessment of
disentanglement, we employ two metrics: a prediction-based indicator:
Disentanglement, Completeness, and Informativeness (DCI) [11], and an
intervention-based criterion: Total AUROC Difference (TAD) [52]. Meanwhile, we
report the reconstruction quality in Appendix G and conclude that ParamReL can
achieve near-exact reconstruction on Celeba (Figure 12), 3DShapes (Figure 13)
and CIFAR10 (Figure 14).

Table 3: Comparison of representation learning algorithms on continuous data
by disentanglement performance (mean ± std) and classification. The
quantitative results for each algorithm averaged over five trials.
($\mathit{g}$ for the Gaussian distribution, $\mathit{c}$ for the categorical
distribution, and $\mathit{d}$ for the delta distribution; $\uparrow$ higher
is better; [Top-1,Top-2]).

|  |  | CelebA | 3DShapes
---|---|---|---|---
| Prior |  |  TAD$\uparrow$ |  ATTRS$\uparrow$ |  NotBlur? |  AUROC$\uparrow$ |  DCI$\uparrow$ |  NotBlur?
S | - |  AE |  0.042±0.004 |  1.0±0.0 |  ✖ |  0.759±0.003 |  0.219±0.001 |  ✖
$\mathit{g}$ |  VAE |  0.000±0.000 |  0.0±0.0 |  ✖ |  0.770±0.002 |  0.276±0.001 |  ✖
$\mathit{g}$ |  $\beta$-VAE |  0.088±0.051 |  1.6±0.8 |  ✖ |  0.699±0.001 |  0.281±0.001 |  ✖
$\mathit{g}$ |  infoVAE |  0.000±0.000 |  0.0±0.0 |  ✖ |  0.757±0.003 |  0.134±0.001 |  ✖
$\mathit{g}$ |  DiffAE |  0.155±0.010 |  2.0±0.0 |  ✔ |  0.799±0.002 |  0.196±0.001 |  ✔
$\mathit{g}$ |  infoDiffusion |  0.299±0.006 |  3.0±0.0 |  ✔ |  0.848±0.001 |  0.342±0.002 |  ✔
P | $\mathit{c}$ |  ParamReL($\lambda=0.1$) |  0.261±0.01 |  5.0±0.0 |  ✔ |  0.846±0.349 |  0.477±0.052 |  ✔
P | $\mathit{d}$ |  ParamReL($\lambda=0.01$) |  0.302±0.005 |  4.0±0.0 |  ✔ |  0.850±0.116 |  0.567±0.005 |  ✔
$\mathit{d}$ |  ParamReL($\lambda=0.1$) |  0.368±0.005 |  3.0±0.0 |  ✔ |  0.865±0.064 |  0.485±0.041 |  ✔

The qualitative latent traversal and quantitative disentanglement metrics show
that learning with ParamReL leads to a visual traversal, which intuitively
matches the attribute on which the latent representation is an effective
detector.

##  6 Conclusion and Limitations

A novel parameter space representation learning framework ParamReL is
introduced in this work, which provides a uniform strategy for dealing with
continuous, discretized, and discrete data. Unlike previous encoder methods
that encode observations into static latent semantics, a self-encoder is
capable of obtaining a series of progressively structured latent semantics
from the step-wise parameters. Experimental results on tasks of latent
interpolation, disentanglement, time-varying conditional reconstruction, and
conditional generation validate the merits of ParamReL, demonstrating its
superior performance in gaining unified representations and a clearer visual
understanding.

ParamReL discloses interesting follow-ups, including: (1) The precision
variables $\alpha_{t}$ could potentially be optimized to accelerate the
sampling procedure, which will be further explored. (2) Utilizing a pre-
trained model for the U-Net architecture might enhance the performance of the
proposed ParamReL, which will be investigated in future work.

## References

  * [1] Alexander Alemi, Ben Poole, Ian Fischer, Joshua Dillon, Rif A Saurous, and Kevin Murphy.  Fixing a broken elbo.  In Proceedings of the International Conference on Machine Learning, pages 159–168. PMLR, 2018.
  * [2] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg.  Structured denoising diffusion models in discrete state-spaces.  Advances in Neural Information Processing Systems, 34:17981–17993, 2021.
  * [3] Juhan Bae, Michael R Zhang, Michael Ruan, Eric Wang, So Hasegawa, Jimmy Ba, and Roger Grosse.  Multi-rate vae: Train once, get the full rate-distortion curve.  arXiv preprint arXiv:2212.03905, 2022.
  * [4] Yoshua Bengio, Aaron Courville, and Pascal Vincent.  Representation learning: A review and new perspectives.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 35(8):1798–1828, 2013.
  * [5] Chris Burgess and Hyunjik Kim.  3d shapes dataset.  https://github.com/deepmind/3dshapes-dataset/, 2018.
  * [6] Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud.  Isolating sources of disentanglement in variational autoencoders.  Advances in Neural Information Processing Systems, 31:2615–2625, 2018.
  * [7] Ting Chen, Ruixiang Zhang, and Geoffrey Hinton.  Analog bits: Generating discrete data using diffusion models with self-conditioning.  arXiv preprint arXiv:2208.04202, 2022.
  * [8] Li Deng.  The mnist database of handwritten digit images for machine learning research.  IEEE Signal Processing Magazine, 29(6):141–142, 2012.
  * [9] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  Advances in Neural Information Processing Systems, 34:8780–8794, 2021.
  * [10] Sixun Dong, Huazhang Hu, Dongze Lian, Weixin Luo, Yicheng Qian, and Shenghua Gao.  Weakly supervised video representation learning with unaligned text for sequential videos.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2437–2447. IEEE, 2023.
  * [11] Cian Eastwood and Christopher KI Williams.  A framework for the quantitative evaluation of disentangled representations.  In Proceedings of the International Conference on Learning Representations, 2018.
  * [12] Babak Esmaeili, Robin Walters, Heiko Zimmermann, and Jan-Willem van de Meent.  Topological obstructions and how to avoid them.  Advances in Neural Information Processing Systems, 36, 2023.
  * [13] B Everett.  An introduction to latent variable models.  Springer Science & Business Media, 2013.
  * [14] Jean-Yves Franceschi, Aymeric Dieuleveut, and Martin Jaggi.  Unsupervised scalable representation learning for multivariate time series.  Advances in Neural Information Processing Systems, 32, 2019.
  * [15] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial nets.  Advances in Neural Information Processing Systems, 27, 2014.
  * [16] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez.  Bayesian flow networks.  arXiv preprint arXiv:2308.07037, 2023.
  * [17] Irina Higgins, Loic Matthey, Arka Pal, Christopher P Burgess, Xavier Glorot, Matthew M Botvinick, Shakir Mohamed, and Alexander Lerchner.  beta-vae: Learning basic visual concepts with a constrained variational framework.  Proceedings of the International Conference on Learning Representations, 3, 2017.
  * [18] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
  * [19] Geonho Hwang, Jaewoong Choi, Hyunsoo Cho, and Myungjoo Kang.  Maganet: Achieving combinatorial generalization by modeling a group action.  In Proceedings of the International Conference on Machine Learning, pages 14237–14248. PMLR, 2023.
  * [20] Zhuxi Jiang, Yin Zheng, Huachun Tan, Bangsheng Tang, and Hanning Zhou.  Variational deep embedding: An unsupervised and generative approach to clustering.  arXiv preprint arXiv:1611.05148, 2016.
  * [21] Hyunjik Kim and Andriy Mnih.  Disentangling by factorising.  In Proceedings of the International Conference on Machine Learning, pages 2649–2658. PMLR, 2018.
  * [22] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv preprint arXiv:1312.6114, 2013.
  * [23] Rahul Krishnan, Dawen Liang, and Matthew Hoffman.  On the challenges of learning with inference networks on sparse, high-dimensional data.  In Proceedings of International Conference on Artificial Intelligence and Statistics, pages 143–151. PMLR, 2018.
  * [24] Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  Technical report, University of Toronto, 2009.
  * [25] James Kwok and Ryan P Adams.  Priors for diversity in generative latent variable models.  Advances in Neural Information Processing Systems, 25:3005–3013, 2012.
  * [26] Tianhong Li, Huiwen Chang, Shlok Kumar Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan.  Mage: Masked generative encoder to unify representation learning and image synthesis.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2142–2152. IEEE, 2023.
  * [27] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.  Deep learning face attributes in the wild.  In Proceedings of the IEEE International Conference on Computer Vision, pages 3730–3738, 2015.
  * [28] Christos Louizos, Kevin Swersky, Yujia Li, Max Welling, and Richard Zemel.  The variational fair autoencoder.  arXiv preprint arXiv:1511.00830, 2015.
  * [29] Cristian Meo, Louis Mahon, Anirudh Goyal, and Justin Dauwels.  $\backslash$alpha TC-VAE: On the relationship between Disentanglement and Diversity.  In Proceedings of the International Conference on Learning Representations, 2024.
  * [30] Khalid Oublal, Said Ladjal, David Benhaiem, Emmanuel LE BORGNE, and François Roueff.  Disentangling time series representations via contrastive independence-of-support on l-variational inference.  In Proceedings of the International Conference on Learning Representations, 2024.
  * [31] Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn.  Diffusion autoencoders: Toward a meaningful and decodable representation.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10619–10629, 2022.
  * [32] Ali Razavi, Aaron Van den Oord, and Oriol Vinyals.  Generating diverse high-fidelity images with vq-vae-2.  Advances in Neural Information Processing Systems, 32:14837–14847, 2019.
  * [33] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Proceedings of the Medical Image Computing and Computer-Assisted Intervention, pages 234–241. Springer, 2015.
  * [34] Huajie Shao, Shuochao Yao, Dachun Sun, Aston Zhang, Shengzhong Liu, Dongxin Liu, Jun Wang, and Tarek Abdelzaher.  Controlvae: Controllable variational autoencoder.  In Proceedings of the International Conference on Machine Learning, pages 8655–8664. PMLR, 2020.
  * [35] Ken Shoemake.  Animating rotation with quaternion curves.  In Proceedings of the Annual Conference on Computer Graphics and Interactive Techniques, pages 245–254, 1985.
  * [36] Casper Kaae Sønderby, Tapani Raiko, Lars Maaløe, Søren Kaae Sønderby, and Ole Winther.  Ladder variational autoencoders.  Advances in Neural Information Processing Systems, 29:3738–3746, 2016.
  * [37] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  arXiv preprint arXiv:2010.02502, 2020.
  * [38] Yuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Wei-Ying Ma.  Unified generative modeling of 3d molecules with bayesian flow networks.  In Proceedings of the International Conference on Learning Representations, 2024.
  * [39] Hiroshi Takahashi, Tomoharu Iwata, Atsutoshi Kumagai, Sekitoshi Kanai, Masanori Yamada, Yuuki Yamanaka, and Hisashi Kashima.  Learning optimal priors for task-invariant representations in variational autoencoders.  In Proceedings of the ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 1739–1748, 2022.
  * [40] Jakub Tomczak and Max Welling.  Vae with a vampprior.  In Proceedings of the International Conference on Artificial Intelligence and Statistics, pages 1214–1223. PMLR, 2018.
  * [41] Sana Tonekaboni, Chun-Liang Li, Sercan O Arik, Anna Goldenberg, and Tomas Pfister.  Decoupling local and global representations of time series.  In Proceedings of the International Conference on Artificial Intelligence and Statistics, pages 8700–8714. PMLR, 2022.
  * [42] Michael Tschannen, Olivier Bachem, and Mario Lucic.  Recent advances in autoencoder-based representation learning.  arXiv preprint arXiv:1812.05069, 2018.
  * [43] Arash Vahdat and Jan Kautz.  Nvae: A deep hierarchical variational autoencoder.  Advances in Neural Information Processing Systems, 33:19667–19679, 2020.
  * [44] Aaron Van Den Oord, Oriol Vinyals, et al.  Neural discrete representation learning.  Advances in Neural Information Processing Systems, 30:6306–6315, 2017.
  * [45] Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov.  Infodiffusion: Representation learning using information maximizing diffusion models.  arXiv preprint arXiv:2306.08757, 2023.
  * [46] Zhangkai Wu, Longbing Cao, and Lei Qi.  evae: Evolutionary variational autoencoder.  IEEE Transactions on Neural Networks and Learning Systems, 2024.
  * [47] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  arXiv preprint arXiv:1708.07747, 2017.
  * [48] Tim Z Xiao and Robert Bamler.  Trading information between latents in hierarchical variational autoencoders.  arXiv preprint arXiv:2302.04855, 2023.
  * [49] Jie Xu, Yazhou Ren, Huayi Tang, Xiaorong Pu, Xiaofeng Zhu, Ming Zeng, and Lifang He.  Multi-vae: Learning disentangled view-common and view-peculiar visual representations for multi-view clustering.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9234–9243, 2021.
  * [50] Kaiwen Xue, Yuhao Zhou, Shen Nie, Xu Min, Xiaolu Zhang, Jun Zhou, and Chongxuan Li.  Unifying bayesian flow networks and diffusion models through stochastic differential equations.  arXiv preprint arXiv:2404.15766, 2024.
  * [51] Tao Yang, Xuanchi Ren, Yuwang Wang, Wenjun Zeng, and Nanning Zheng.  Towards building a group-based unsupervised representation disentanglement framework.  arXiv preprint arXiv:2102.10303, 2021.
  * [52] Eric Yeats, Frank Liu, David Womble, and Hai Li.  Nashae: Disentangling representations through adversarial covariance minimization.  In Proceedings of the European Conference on Computer Vision, pages 36–51. Springer, 2022.
  * [53] Haojie Zhao, Dong Wang, and Huchuan Lu.  Representation learning for visual object tracking by masked appearance transfer.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 18696–18705. IEEE, 2023.
  * [54] He Zhao, Piyush Rai, Lan Du, Wray Buntine, Dinh Phung, and Mingyuan Zhou.  Variational autoencoders for sparse and overdispersed discrete data.  In Proceedings of the International Conference on Artificial Intelligence and Statistics, pages 1684–1694. PMLR, 2020.
  * [55] Shengjia Zhao, Jiaming Song, and Stefano Ermon.  Infovae: Information maximizing variational autoencoders.  arXiv preprint arXiv:1706.02262, 2017.
  * [56] Shengjia Zhao, Jiaming Song, and Stefano Ermon.  Infovae: Balancing learning and inference in variational autoencoders.  In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 5885–5892, 2019.
  * [57] Mingyuan Zhou, Tianqi Chen, Zhendong Wang, and Huangjie Zheng.  Beta diffusion.  Advances in Neural Information Processing Systems, 36, 2023.

##  Appendix A Preliminaries

###  A.1 Bayesian Flow Distribution

Bayesian flow distribution
$p_{\mathrm{F}}(\cdot\mid{\boldsymbol{\mathbf{x}}};t)$ is the marginal
distribution over input parameters at time $t$, given prior distribution,
accuracy schedule $\alpha$ and Bayesian update distribution
$p_{U}(\cdot\mid{\boldsymbol{\mathbf{\theta}}},{\boldsymbol{\mathbf{x}}};\alpha)$,
as follows:

$p_{\mathrm{F}}({\boldsymbol{\mathbf{\theta}}}\mid{\boldsymbol{\mathbf{x}}};t)=p_{U}\left({\boldsymbol{\mathbf{\theta}}}\mid{\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{x}}};\beta(t)\right).$
(8)

###  A.2 Generative Latent Variable Models for Representation Learning

Latent Variable Models (LVMs) [13] which aim at learning the joint
distribution $p({\boldsymbol{\mathbf{x}}},{{\boldsymbol{\mathbf{z}}}})$ over
data ${\boldsymbol{\mathbf{x}}}$ and latent variables
${{\boldsymbol{\mathbf{z}}}}$ present efficient ways for uncovering hidden
semantics. In LVMs, the joint distribution
$p({\boldsymbol{\mathbf{x}}},{\boldsymbol{\mathbf{z}}})$ is usually decomposed
as
$p({\boldsymbol{\mathbf{x}}},{\boldsymbol{\mathbf{z}}})=p({\boldsymbol{\mathbf{x}}}\mid{{\boldsymbol{\mathbf{z}}}})p({{\boldsymbol{\mathbf{z}}}})$,
where $p({{\boldsymbol{\mathbf{z}}}})$ represents prior knowledge for
inference [42], thus facilitating learning the conditional distribution
$p({\boldsymbol{\mathbf{x}}}\mid{{\boldsymbol{\mathbf{z}}}})$. Among LVMs,
Variational AutoEncoders (VAEs) [22] and diffusion models [18, 37] are two
representative approaches [25].

In VAEs, latent variables ${\boldsymbol{\mathbf{z}}}$ is obtained through an
_encoder network_
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}\,|\,{\boldsymbol{\mathbf{x}}})$,
whereas observations are reconstructed through a _decoder network_
$p_{{\boldsymbol{\mathbf{\theta}}}}({\boldsymbol{\mathbf{x}}}\,|\,{\boldsymbol{\mathbf{z}}})$,
with ${\boldsymbol{\mathbf{\phi}}}$ and ${\boldsymbol{\mathbf{\theta}}}$ being
the encoder and decoder parameters. The dimensions of
${\boldsymbol{\mathbf{z}}}$ are usually much smaller than those of
${\boldsymbol{\mathbf{x}}}$, denoted as
$|{\boldsymbol{\mathbf{z}}}|\ll|{\boldsymbol{\mathbf{x}}}|$, such that
redundant information is effectively removed and the most semantically
meaningful factors are abstracted [28]. VAEs are popular for downstream tasks
like disentanglement [17, 51, 19, 12], classification [39, 41], and clustering
[20, 49].

On the other hand, diffusion models [18, 37] first use $T$ diffusion steps to
transform observation ${\boldsymbol{\mathbf{x}}}$ into a white noise
${\boldsymbol{\mathbf{x}}}_{T}$ and then use $T$ denoising steps to
reconstruct the observation. Diffusion models have obtained impressive
performance in the fidelity and diversity of generation tasks. However, they
might be unable to obtain meaningful latent semantics since the dimensions of
${\boldsymbol{\mathbf{x}}}$ and ${\boldsymbol{\mathbf{x}}}_{T}$ are the same
as $|{\boldsymbol{\mathbf{x}}}|=|{\boldsymbol{\mathbf{x}}}_{T}|$. [31, 45]
have attempted to integrate a decodable auxiliary variable
${\boldsymbol{\mathbf{z}}}$ to enable diffusion models to obtain low-
dimensional latent semantics. However, they have not overcome issues like the
slow training speed inherent to the diffusion and reverse processes.

###  A.3 Illustration of Parameter Space Optimization

Figure 6 illustrates the optimal data distribution learned in the parameter
space. The plot presents stochastic parameter trajectories for the input
distribution mean (indicated by white lines) overlaid on a Bayesian flow
distribution logarithmic heatmap.

Figure 6: This figure illustrates optimization in the parameter space after
$t$ iterations.

##  Appendix B ELBO of ParamReL

We derive the ELBO of ParamReL defined in Eq. (3).

$\displaystyle\quad\log p({\boldsymbol{\mathbf{x}}}_{0})$
$\displaystyle=\log\int_{\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}}\int_{\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}}p\left({{\boldsymbol{\mathbf{x}}}}_{0},\\{{{\boldsymbol{\mathbf{x}}}}_{t}\\}_{t},\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}\,|\,{\boldsymbol{\mathbf{\theta}}}_{0},\alpha\right)\mathrm{d}\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}\mathrm{d}\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}$
$\displaystyle=\log\int_{\\{{\boldsymbol{\mathbf{\theta}}}_{t}\\}_{t}}\int_{\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}}\int_{\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}}p(\\{{\boldsymbol{\mathbf{\theta}}}_{t}\\}_{t}|-)p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))\prod_{t=T}^{1}p({\boldsymbol{\mathbf{z}}}_{t})\mathbb{E}_{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}))}[p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})]$
$\displaystyle\qquad\qquad\mathrm{d}\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}\mathrm{d}\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}\mathrm{d}\\{{\boldsymbol{\mathbf{\theta}}}_{t}\\}_{t}$
$\displaystyle=\log\int_{\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}}\int_{\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}}\int_{\\{{\boldsymbol{\mathbf{\theta}}}_{t}\\}_{t}}p(\\{{\boldsymbol{\mathbf{\theta}}}_{t}\\}_{t}|-)\frac{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))\prod_{t=T}^{1}p({\boldsymbol{\mathbf{z}}}_{t})\mathbb{E}_{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}))}[p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})]}{\prod_{t=1}^{T}p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)}$
$\displaystyle\qquad\qquad\qquad\cdot\prod_{t=1}^{T}p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)\mathrm{d}\\{{\boldsymbol{\mathbf{z}}}_{t}\\}_{t}\mathrm{d}\\{{\boldsymbol{\mathbf{x}}}_{t}\\}_{t}\mathrm{d}\\{{\boldsymbol{\mathbf{\theta}}}_{t}\\}_{t}$
$\displaystyle\geq\mathbb{E}_{\prod_{t=1}^{T}p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)p({\boldsymbol{\mathbf{\theta}}}_{t}|-)}\left[\log\frac{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))\prod_{t=T}^{1}p({\boldsymbol{\mathbf{z}}}_{t})\mathbb{E}_{p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}))}[p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})]}{\prod_{t=1}^{T}p_{\mathrm{S}}({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{\boldsymbol{\mathbf{x}}}_{t};\alpha_{t})q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)}\right]$
$\displaystyle=\sum_{t=1}^{T}\mathbb{E}_{p_{F}({\boldsymbol{\mathbf{\theta}}}_{t}|-)}\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t})}\left\\{\mathbb{E}_{p_{\mathrm{S}}\left({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{{\boldsymbol{\mathbf{x}}}_{0}};\alpha_{T:t}\right)}\left[\log\frac{p_{\mathrm{S}}\left({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{{\boldsymbol{\mathbf{x}}}_{0}};\alpha_{T:t}\right)}{p_{\mathrm{R}}\left({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}),\alpha_{t}\right)}\right]\right.$
$\displaystyle\qquad\qquad\qquad\left.-\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}\left({{\boldsymbol{\mathbf{z}}}}_{t}\,|\,{{\boldsymbol{\mathbf{\theta}}}_{t}}\right)}\left[\log\frac{q_{{\boldsymbol{\mathbf{\phi}}}}\left({{\boldsymbol{\mathbf{z}}}}_{t}\,|\,{{\boldsymbol{\mathbf{\theta}}}_{t}}\right)}{p({{\boldsymbol{\mathbf{z}}}}_{t})}\right]\right\\}+\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{0},{\boldsymbol{\mathbf{\theta}}}_{0})}\left[{\ln
p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))}\right]$
$\displaystyle=-\sum_{t=1}^{T}\mathbb{E}_{p_{F}({\boldsymbol{\mathbf{\theta}}}_{t}|-)}\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t})}\left\\{{\textsc{kl}\left[p_{\mathrm{S}}\left({\boldsymbol{\mathbf{x}}}_{t-1}\,|\,{{\boldsymbol{\mathbf{x}}}_{0}};\alpha_{T:t}\right)\;\|\;p_{\mathrm{R}}\left({\boldsymbol{\mathbf{x}}}_{t-1};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}),\alpha_{t}\right)\right]}\right.$
$\displaystyle\qquad\qquad\qquad\left.-\textsc{kl}\left[q_{{\boldsymbol{\mathbf{\phi}}}}\left({{\boldsymbol{\mathbf{z}}}}_{t}\,|\,{{\boldsymbol{\mathbf{\theta}}}_{t}}\right)\;\|\;p({{\boldsymbol{\mathbf{z}}}}_{t})\right]\right\\}+\mathbb{E}_{q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{0},{\boldsymbol{\mathbf{\theta}}}_{0})}\left[{\ln
p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{0};\psi({\boldsymbol{\mathbf{\theta}}}_{0},{\boldsymbol{\mathbf{z}}}_{0}))}\right]:=\mathcal{L}_{\texttt{ParamReL}}$
(9)

##  Appendix C Illustration of Network Architecture

###  C.1 Latent Encoding

We define the encoder
$q_{{\boldsymbol{\mathbf{\phi}}}}({\boldsymbol{\mathbf{z}}}_{t}|{\boldsymbol{\mathbf{\theta}}}_{t},t)$
with a latent encoding network parameterized by a U-Net, and

the decoder as:

$p_{\mathrm{O}}({\boldsymbol{\mathbf{x}}}_{t};\psi({\boldsymbol{\mathbf{\theta}}}_{t},{\boldsymbol{\mathbf{z}}}_{t}))=\delta\left({{\boldsymbol{\mathbf{x}}}}-\frac{\boldsymbol{\mu}}{\gamma(t)}-\sqrt{\frac{1-\gamma(t)}{\gamma(t)}}\epsilon_{\psi}({{\boldsymbol{\mathbf{\theta}}}}_{t},{\boldsymbol{\mathbf{z}}}_{t})\right),$
(10)

with a noise prediction network
$\epsilon_{\psi}({{\boldsymbol{\mathbf{\theta}}}}_{t},{\boldsymbol{\mathbf{z}}}_{t})$
parameterized by a U-Net [33] same with [18]. Following the previous work [9],
we employ an adaptive group normalization layer (AdaGN) to do conditional
embedding, incorporating the input parameters and timestep into each residual
block after a group normalization:

$\operatorname{AGN}({\boldsymbol{\mathbf{\theta}}}_{t},t)=(1+\mathbf{s}(t))\cdot\operatorname{GroupNorm}({\boldsymbol{\mathbf{\theta}}}_{t})+\mathbf{b}(t),$
(11)

where the $\mathbf{s}(\cdot)$ function outputs scaling factors based on the
conditional input $t$, dynamically adjusting the normalized features to adapt
to different contexts, and the $\mathbf{b}(\cdot)$ function provides bias
terms derived from $t$, allowing shifting of the normalized features according
to the given conditions. In decoders, the AGN layer will be as follows:

$\operatorname{AGN}(\mathbf{\theta}_{t},\mathbf{z}_{t},t)=\left(1+\mathbf{s}(t)\right)\cdot\operatorname{GroupNorm}\left(\left(1+\mathbf{s}(\mathbf{z}_{t})\right)\cdot\operatorname{GroupNorm}(\mathbf{\theta}_{t})+\mathbf{b}(\mathbf{z}_{t})\right)+\mathbf{b}(t).$
(12)

###  C.2 Modified U-Net

Similar to the diffusion-based representation learning model, we update the
U-Net architecture based on Residual Blocks and Attention Modules. However,
unlike previous approaches [18, 37, 31, 45], we use shallower layers in the
upper and down modules while incorporating an additional attention mechanism
in the bottleneck module to achieve significant representations. Figure 7
illustrates the specific structural differences.

Figure 7: U-Net comparisons of infoDiffusion (a) and ours (b). We apply the
Attention module in the bottleneck layer, shallower than the infoDiffusion’s
U-Net.

###  C.3 Encoder

We apply the same U-Net frameworks to encoders. The framework of encoders is
depicted in Figure 8.

Figure 8: Two kinds of Encoders in ParamReL, including vanilla encoders (a)
and progressive encoders (b) conditioned on time $t$.

The sizes of the model parameters for each configuration are provided in Table
4.

Table 4: Comparisons of parameters in training. We report the parameters of infoDiffusion and ParamReL on the CelebA dataset with $(3\times 64\times 64)$ input size.  |  | Encoder | Decoder
---|---|---|---
CelebA | infoDiffusion | 143.89MB | 200.18MB
$(3\times 64\times 64)$ | ParamReL | 245.41MB | 324.31MB

##  Appendix D Hyperparameters for Training

Table 5 presents the hyperparameter settings for training ParamReL. Different
bin values are provided for various continuous datasets. All models are
trained for 50 epochs. “Channel mult" denotes the channel shapes in each
ResNet block within the U-Net architecture.

Table 5: Hyperparameters for training Bayesian Flow Networks, U-Net architecture, training protocols, and devices. | Discrete | Continuous
---|---|---
|  | bFashionMNIST | bMNIST | CelebA | 3DShapes | CIFAR10
BFN | bins | - | - | 256 | 256 | 16
Exps | Latent Dim | 32 | 32 | 64 | 32 | 32
Batch Size | 128 | 128 | 64 | 64 | 128
Learning Rate | 1e-4 | 1e-4 | 2e-4 | 2e-4 | 1e-4
DataShape | (28, 28, 1) | (28, 28, 1) | (64,64,3) | (64,64,3) | (32,32,3)
Models | Channel Mult | [1,2,4] | [1,2,4] | [1,2,2,2] | [1,2,2,2] | [1,2,4]
Num. Channels | 32 | 32 | 64 | 64 | 32
Device | GPU | H100 | H100 | H100 | H100 | H100

##  Appendix E Latent Interpolation

The latent space interpolation can be described as follows. Firstly, we noise
source images to generate latent pairs by sender distribution,
$<{\boldsymbol{\mathbf{x}}}^{1}_{1},{\boldsymbol{\mathbf{x}}}^{2}_{1}>$, where
${\boldsymbol{\mathbf{x}}}^{1}_{1}\sim
q(\cdot\mid{\boldsymbol{\mathbf{x}}}^{1}_{N})$ and
${\boldsymbol{\mathbf{x}}}^{2}_{1}\sim
q(\cdot\mid{\boldsymbol{\mathbf{x}}}^{2}_{N})$. Then, we implement two methods
from [35] to generate four interpolated latent pairs
$\bar{{\boldsymbol{\mathbf{x}}}}_{1:4}$, i.e., linear interpolation, and
spherical interpolation:

$\displaystyle\bar{{\boldsymbol{\mathbf{x}}}}_{i}$
$\displaystyle=(1-\lambda){\boldsymbol{\mathbf{x}}}_{0}^{1}+\lambda{\boldsymbol{\mathbf{x}}}_{0}^{2},$
(13) $\displaystyle\bar{{\boldsymbol{\mathbf{x}}}}_{i}$
$\displaystyle=\frac{\sin((1-\alpha)\theta)}{\sin(\theta)}{\boldsymbol{\mathbf{x}}}_{0}^{1}+\frac{\sin(\alpha\theta)}{\sin(\theta)}{\boldsymbol{\mathbf{x}}}_{0}^{1},$

where $\lambda$ is the scale coefficient, $\alpha\in[0,1]$ denotes the
interpolation steps, and
$\theta=\arccos\left(\frac{\left({\boldsymbol{\mathbf{x}}}_{0}^{1}\right)^{\top}{\boldsymbol{\mathbf{x}}}_{0}^{2}}{\left\|{\boldsymbol{\mathbf{x}}}_{0}^{1}\right\|\left\|{\boldsymbol{\mathbf{x}}}_{0}^{2}\right\|}\right)$
is the angle between ${\boldsymbol{\mathbf{x}}}_{0}^{1}$ and
${\boldsymbol{\mathbf{x}}}_{0}^{2}$.

Figure 9: Samples reconstructed from our trained ParamReL on dataset Binary-
MNIST. Figure 10: Samples generated from our trained ParamReL. Figure 11:
Details of samples generated from our trained ParamReL.

##  Appendix F Generation Quality

We illustrate the unconditional generation quality in Figure 10 and Figure 11
on bMNIST. Images sampled from VAE-based model are blurry, as shown in Figure
10 (b). We implement two sampling strategies in the Diffusion-based model
[45], and both can only sample grey-scale images. Figure 10 (c) is sampled
from the DDIM sampler, and Figure 10 (d) is sampled from a two-phased sampling
procedure: form timesteps $T$ to $T/2$, denoise and sample using a pre-trained
vanilla denoising diffusion model. For timesteps ranging from $T/2$ to $0$,
proceed with sampling utilizing the InfoDiffusion model. Figure 10 (e) is
images generated from our ParamReLc model. Refer to Figure 11 for more
detailed information per image sampled from corresponding models. We can
conclude that ParamReL can be sampled from the discrete distribution where the
image value is binarized.

Figure 12: Samples reconstructed from our trained ParamReL on CelebA.

##  Appendix G Reconstruction Quality

The qualitative reconstruction results are shown in Figure 12 for CelebA,
Figure 13 for 3DShapes and Figure 13 for CIFAR10.

Figure 13: Samples reconstructed from our trained ParamReL on 3DShapes. Figure
14: Samples reconstructed from our trained ParamReL on CIFAR10.