Trans-Dimensional Generative Modeling via Jump Diffusion Models

Summary

This paper proposes a new diffusion model based on jump diffusion processes. Compared with previous discrete and continuous formulations, the model introduces the usage of a transition kernel, which models the jump process in a semantically meaningful manner. The method absorbs standard constructions like diffusion guidance and produces good results on molecule and robot arm video tasks.

Strengths

* The proposed methodology is elegant and theoretically sound. * Many constructions like diffusion guidance extend naturally. * The proposed methodology is widely applicable, as many natural data types include a continuous portion (modeling the space) as well as a discrete but unknown component.

Weaknesses

* For the molecule task, the metrics are based on molecular validity properties. This seems like it wouldn't account for overfitting from the model, which seems more relevant for a generative modeling task. In particular, can you report a diversity/non-memorization metric? * The robot arm example, while sufficient for the paper, is rather toy.

Questions

Nothing more than the weaknesses section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

This paper focuses on varying dimensional datasets and proposes a novel generative model to solve the varying dimensional problems. The proposed model is theoretically valid and has an interesting and novel contribution to extending the traditional score-based generative model by generating both state values and dimensions jointly during the generative process, which is idea-simple but effective. Experiments on molecule generation and video generation both show this model's effectiveness.

Strengths

1. The idea of jointly modeling the state and dimension, in particular, the idea of using the intensity function to model the jump distribution, is interesting and novel. 2. The theoretical contribution is valid and the experiment is thorough in supporting the proposed model.

Weaknesses

None.

Questions

1. I am confused about how to properly define $K^{{del}}(i|n)$ in $\overrightarrow{K}_t(\mathbf{Y}|\mathbf{X})$? Any clarification about this issue will be helpful.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes, the authors have adequately addressed the limitations and potential negative societal impact of their work.

Summary

This paper addresses the problem of modelling data of various dimensions. This is achieved by generalising diffusion models as jump diffusion processes, allowing the content and dimension of data to be jointly modelled. The forward process gradually corrupts the data with gaussian noise while also gradually deleting dimensions until a single normally distributed dimension remains. The reverse process learns both the score function, when to add a dimension, and what data to place in the newly created dimension. Experimental results on molecule and video generation tasks show that the approach can well represent the distribution of data dimensions, while outperforming/being competitive in terms of sample quality.

Strengths

- The proposed approach of modelling the generative process as a jump diffusion process is interesting and is a sensible solution. - The paper is very clear, easy to read, and to my understanding, the method is technically sound. I particularly like the intuitive explanation of the loss in Equations 3-4; and predicting the original data dimension (line 191) is a good idea to address the optimisation issues. - Experimental results demonstrate the effectiveness of the proposed approach (Table 2), performing comparability to, or outperforming the baseline which sample dimension prior to sampling. I also like the evaluation of the impact of setting $\lambda=0$ (dimension deletion/insertion rate) towards the end of the diffusion process. - It is shown that reconstruction guidance can be used to generate molecules with specified features (Table 3) much better than dimension independent approaches. - The problem of jointly modelling content and dimension of data is an important one to address. In my opinion, the proposed solution is compelling and I believe it will be very useful to others.

Weaknesses

- Predicting the content of the newly created dimensions with a gaussian distribution is limiting, potentially creating a discrepancy between train/test time and requiring more diffusion time to correct. - Fixing $\lambda=0$ towards the end of the diffusion process to ensure added dimensions have time to diffuse is a weakness of the proposed approach. Particularly for more complex data, there is no guarantee that the remaining time will be sufficient. - Finally, it would have been nice to see a perceptual quality metric for the video generation task, this would offset the above weaknesses.

Questions

- Do the authors have any rationale/evidence that fixing $\lambda$ at $t<0.1$ is a reasonable choice? For instance, graphing metrics over a variety of limits. - Similarly, is there rationale why predicting new dimensions with a gaussian is sufficient; or if for more complex data this is problematic, is there a more expressive extension?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Limitations are well discussed throughout the paper. The previously discussed weaknesses could be mentioned as well.

Summary

This paper proposes jump diffusion, which is a novel diffusion model to handle data with varying dimensions. The proposed method is derived from a special forward process that contains a jump part that changes the dimension of the generated samples. The corresponding backward process and the learning objective are derived, with two more components that need to be learned other than the standard diffusion drift. Numerical issues are properly handled to make the mathematical model work in practice. Experiments over various application scenarios are conducted to demonstrate the flexibility and versatility of the proposed framework.

Strengths

1. The proposed method targets at a very important problem. It is a great contribution to invent theoretically sound diffusion models to deal with data with varying dimensions. This kind of problem frequently happens in real-world applications. 2. The motivation is strong, the idea is reasonable, and the mathematical derivations makes the idea a solid framework. 3. Applications on molecule and video generation showcases the potential of the proposed framework in reality.

Weaknesses

1. Although I do think the proposed framework is interesting, generation with varying dimensionality can be addressed easily by (1) learning a distribution over the dimension numbers and sample dimension number from the learned distribution (2) sampling initial random noises according to the sampled dimension number at the beginning of the generation and pad the unused dimensions with 0. This simple method has little modification over existing diffusion models, and the only extra efforts one need to do is to add paddings during training. It would be very nice if I can hear comments from the authors on this simple baseline. I do appreciate the theoretical contribution of this paper, and it would be better if the authors can compare with this simple baseline . 2.In L84, should the summation be taken over m<n? Because there is only deletion in the forward process.

Questions

I have one question on the proposed diffusion model and I would appreciate answers from the authors. In Algorithm Box 1, what is the initial distribution? Is it a random standard Gaussian noise with only 1 dimension? If so, how do we decide which dimension is the initial dimension? For example, if we have x = (3, x_1, x_2, x_3), what is the initial random noise? (1, N(0, I), 0, 0), (1, 0, N(0, I), 0), or (1, 0, 0, N(0, I) ? Or are they uniformly sampled? Moreover, if the initial number of dimension is always 1, how do we guarantee that in setting the forward process? I feel like there must be some constraints on the scheduling of rate function $\lambda(t)$. Maybe I missed some parts of the derivation in the paper.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

The limitations and broader impacts are properly discussed.

Response to Authors

Thanks a lot to the authors for their thorough responses. The rebuttal addresses my concerns and I am still of the opinion that this is a strong paper and advocate acceptance. In particular, the provided FID scores added in the rebuttal pdf showing that the final video frames to be generated are of comparable quality, is compelling evidence that the Gaussian approximation is sufficient in this setting. It is also true that should it prove to be problematic, a more expressive model could be used. I also appreciate the extra experiment on $\lambda_t$ showing that $0.1T$ is a reasonable choice and the impact on sample properties when this value is changed, and agree that this is a reasonable additional hyperparameter to have.

Thank you

Thank you for your comments comparing between your model and my simple baseline! The joint modeling of both dimensionality and state is indeed important. I agree. Other replies also clearly answer my questions. Although the initial distribution in the backward and forward process has a little bit mismatch, I think the way the authors handle it can mitigate the influence. Thank you for the reply again!