Parameter Space Representation Learning on Mixed-type Data

• In lines 236-237 the authors mention “Given the straightforward definition of Bayesian update function h(·), its inverse operation is generally easy to derive. The details of such results can be found in Figure 14.” However, Figure 14 does not provide any details on the Bayesian update function. Were the authors referring to Table 2?
• Have the authors tried to optimize the ELBO directly instead of using an additional MI term? It would be interesting to see where the former goes wrong and to justify why the MI term is actually needed.
• In Section 3.5, the caption is called “Variational Inference for Intractable Joint Distribution”. However, the joint distribution seems to be tractable, the authors even give an expression.
• How long does it take to train the method in comparison to other methods?
• What does the ATTRS column in Table 1 indicate? Is that mentioned somewhere?
• What is the influence of $\gamma,\lambda$? What happens if I set them low very low/high?

1 Introduction

• The authors correctly cite prior work highlighting the challenges of mixed-type data representation learning. However, explicitly mentioning real-world applications that use mixed-type data latent representations could clarify the motivation of this work. Additionally, it would help to articulate the specific advantages of parameter-space representation learning over output-space representation. (l. 37-42)

• The promise of using Bayesian Flow Networks (BFNs) as a foundation would be more compelling if the authors provided a clearer intuition on why BFNs are better suited for handling mixed-type data than diffusion models and how they refine parametric distributions. Furthermore, what limitations prevent BFNs from capturing low-dimensional latent semantics? How would a latent representation benefit the model: are we aiming to enable latent space operations, similar to those seen in VAEs and GANs, such as latent walks? (l. 43-48)

• The authors refer to “parameters of BFNs” where it might be clearer to specify “parameters of the distribution produced by BFNs” (e.g., line 53). Distinguishing these concepts would add clarity. Additionally, I thought that “parameter space representation learning” referred to learning embeddings of the model’s own parameters, which supports model transferability and interpretability. Since ParamReL instead learns representations of the parameters of the probability distribution generated by BFNs, emphasizing this distinction could help avoid reader confusion.

2 Understanding Bayesian Flow Networks -- An Alternative View

• The first paragraph suggests an “alternative view” of BFNs, yet it’s not immediately clear how this perspective diverges from the original one. Additionally, the point about the accessibility of BFN concepts in the original formulation might not be necessary.

• How do BFNs avoid the expressiveness limitations seen in VAEs, where the variational distribution can yield overly simplistic distributions and result in sample over-smoothing?

3 ParamReL: Parameter Space Representation Learning

• In lines 140-162, it is unclear which specific contributions come from Baranchuk et al. (2021), Rombach et al. (2022), and Luo et al. (2024). Providing clarity on these references would improve the reader’s understanding.

• Some networks are indicated with indexed parameters (e.g., the encoder q_{\theta}), while others, such as \psi, are not. Consistent notation would make it clearer which terms refer to networks.

• I am not familiar with BFNs, but I was wondering why are “the series of latent semantics $\{z_t\}_{t=1}^T$ expected to exhibit progressive semantic changes (e.g., age, smile, skin color)” if they encode the parameters of the data distribution at each time-step (line 177)? Given that the entire distribution is refined over time, what drives the latents across time-steps to model intra-distribution features like age or skin color? Is there a formal justification or an explanation to support this claim?

• Could the authors clarify the notation “p(. |-)” in line 252-253, as "-" is not immediately familiar?

• How does the size of the latent representation $z$ compare to that of the parameters $\theta$?

4 Related Work

• In describing diffusion models, it may be helpful to avoid saying they are “unsuitable” for representation learning without a clear definition of what is meant by “representation learning.” If the authors refer to representation learning as achieving a continuous, lower-dimensional latent space, this distinction should be made clear. Otherwise, the later point about pre-trained diffusion models being used for other tasks, which is related to representation learning, could appear contradictory.

• How does ParamReL differ from InfoDiffusion in terms of learning latent representations?

5 Experiments

• The authors claim to address mixed-type data yet only experiment with discrete and continuous distributions—no mixed-type data is tested, and the only discrete data is binary. To strengthen the contrast with methods like InfoDiffusion, more emphasis on mixed-type data experiments, especially with categorical data, would be helpful.

• In line 330, could the authors define MMD?

1. Diffusion-based representation learning algorithms like InfoDiffusion infer a latent representation from data space directly into latent space. Your proposal infers latent representations from parameter space. Both the data space in InfoDiffusion and the parameter space in your case have the same dimension, don’t they?. What would you say is the reason for the better performance of ParamReL, as compared with InfoDiffusion?
2. Opposite to InfoDiffusion, Rombach et al. [2022], whom the authors cite, trains a diffusion model directly on representation space. Has anyone tried to study the quality of the representations learned by Rombach et al. [2022] from continuous-data as you do in Table 1? What would you say are the key differences between their representations and yours, in the continuous-data case? What makes your approach better?
3. Is the prior distribution over the latent code $p(\mathbf{z}\_t)$ independent of time? In line 171 you write that it "follows a Gaussian distribution", but do the parameters of the Gaussian change with time? If not, why did you decide to constrain the time-dependent posterior with a static prior?
4. Did you include the actual expressions for the inverse of the Bayesian update function $h^{-1}$ somewhere in the paper? I don’t manage to find them.

References

• Wang et al. [2023]: InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models
• Rombach et al. [2022]: High-Resolution Image Synthesis with Latent Diffusion Models

• Why is the evaluation protocol different for continuous and discrete data, and apart from changing the parameter space, are there other differences between the continuous and discrete models?
• Can you elaborate on the significance of the lambda parameter in tables 1 and 4?
• Looking at table 1, you report significantly higher FID scores for DiffAE than was reported in the cited paper. Can you please explain why this is?
• Can you please elaborate on the following statement for figure 4? “(b), the learned semantics exhibit progressive, time-varying changes”