Improving Diffusion-Based Image Synthesis with Context Prediction

Summary

This paper proposes to improve diffusion-based image synthesis by explicitly reinforcing each point to predict its neighborhood context during training, without extra cost at inference. To reduce computation/time complexity of context decoding the authors propose efficient large context decoding adopting Wasserstein distance to characterize the distribution reconstruction loss. The method is applicable to both discrete and continuous diffusion backbones and achieves new SOTA text-to-image generation on MS-COCO with FID 6.21

Strengths

1. The paper is well-written and easy to follow. 2. The method of explicitly reinforcing each point to predict its neighborhood context for diffusion models is well-motivated with effective designs to reduce substantial computation complexity for large-context neighborhood reconstruction. 3. The method is proven effective in boosting FID scores on MS-COCO text-to-image synthesis for both continuous (eDiff-I) and discrete diffusion (VQ-Diffusion) backbones.

Weaknesses

1. The main results are on text-to-image synthesis and image inpainting. It would be good to add unconditional generation results. 2. The method emphasizes on diffusion with better neighbouring context, leading to generations "semantically better consistent with the text prompts"(L233), "prommising cross-modal semantic understanding" (L234), "can synthesize more complex objects and scenes" (L236-237). I don't think the claims are well-justified: e.g. Fig. 2 and Fig. 3 only presents results of the proposed method without any comparison to warrant the aforementioned conclusions. More analysis and evidence of "better semantics" are required other than the overall FID score. 3. Some notations are misleading, e.g. L 114-116, for h_i (h_{t-1}), the subscript is used to indicate both spatial and time; x_t is not defined in the main text.

Questions

1. L260-262: regarding to Fig. 5 I'm not convinced by the observation and conclusion, could the authors make it more clear ? 2. I'm curious about the comparision with Dalle-2 and Imagen in Fig. 7 as the models are not open-sourced.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

No.

Summary

This paper proposes to improve diffusion based image generative training objectives by adding context prediction loss. The motivation of predicting context comes from other non-diffusion based models like semantic segmentation and representation learning. To mitigate the complexity of predicting large per pixel neighbourhood context, the author further models the context as a probability distribution using Wasserstein distance. Experiments show the proposed model achieves new SoTA generation on MSCOCO for both discrete and continuous diffusion models.

Strengths

The introduction is well-written and the motivation of predicting context in the diffusion based model is easy to follow. The presentation of the method including the loss derivation and the training pipeline is easy to understand. The authors conduct intensive experiments showing the proposed context prediction loss can be used on various DM models, achieving SoTA performance on MSCOCO FID and inpainting tasks.

Weaknesses

The paper lacks training and implementation details. For example, the text-to-image experiment uses T5 encoder as the text encoder, but did not mention architecture details and training details. One big motivation to model context as a probability distribution is to improve the training efficiency. As shown in Figure 6, feature matching has lower throughput compared to distribution matching, but it has better FID. I think an important baseline is missing – sampling based feature matching, i.e. use the same number of random samples as proposed in the distribution matching, 9 instead of the full neighbourhoods features for context prediction.

Questions

In Table 1, it is not clear which diffusion model is used with the proposed method for both discrete and continuous diffusion models. For qualitative comparison in Figure 2 and 8, there is no head to head comparison with other baselines, and it is hard to appreciate the improvement.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

There is no limitation/future work discussion in the paper And there is also no training / implementation details in the paper, which cause concerns on reproducibility.

Summary

This paper presents ConPreDiff, a method introduced to improve the performance of diffusion models by preserving the neighborhood context of predicted pixels/features. They achieve this by predicting the neighborhood context during the diffusion generation process. To simplify the modeling complexity, they propose predicting distributions instead of directly reconstructing the neighborhood. The method's effectiveness is demonstrated through extensive experiments on unconditional image generation, text-to-image generation, and image inpainting.

Strengths

* The idea is intuitive and easy to understand. * The proposed method is general and can be easily applied to recent diffusion models. * The performance of the proposed method is very impressive.

Weaknesses

* Recent diffusion models use UNet backbone, which stacks many convolutional and self-attention layers. Thus, it has a large receptive field. Additionally, LDM also has a decoder, which also has a decent receptive field. Therefore, I am confused about the paper's main claim that the point-wise reconstruction neglects to fully preserve the local context. * There are no visual comparisons of the proposed method and baselines. I am not sure if ConPreDiff can really be more local-context consistent compared to other methods. * The authors need to discuss the additional training cost. Besides, they also need to provide the additional parameters they use for the context prediction.

Questions

* L138: Why do you choose $\mathcal P_{\mathcal N_i^s}$ in this form? Is there any insight that motivates you doing this? * L115: What is $\mathbf h_{t-1}$? I suppose it should be $\mathbf x_{t-1}$. * L119: What is $\mathbf h_{t-1}$? I think it should be $\mathbf h_i$. * L138: What is $h_u^{(0)}$? * Figure 4: What is NDM?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

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

The authors did not discuss their limitations and societal impact in the paper.

Summary

This paper proposes an idea of context prediction to boost difussion-based image generation. The core idea is that in each step of diffusion, after the denoised point is generated, neighborhood context prediction is performed. In particular, to maintain the spatial orders of the neighborhood, a permutation invariant loss is used for optimization by replacing the context prediction with neighborhood distribution prediction. Performance improvement against standard diffusion models were presented. in experiments.

Strengths

1. The idea is very interesting and sound. From an image denoising point of view, neighborhood info is commonly used, so it's a natural extension of diffusion-denoising models. 2. Performance improvement showed in experiments are promising.

Weaknesses

The proposed approach probably takes longer to train. Can you discuss from that perspective?

Questions

1. By introducing context, is any sign of "blurriness" introduced? 2. What are the typical cases that do worse compared to standard diffusion?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Similar with text2image papers.

Summary

This paper is proposing context-aware Diffusion Models. They make the models learn the context information by setting up auxiliary networks to estimate the neighbor distributions from the estimated denoised sample from Diffusion Models. The benefit of this approach is that additional cost from the auxiliary networks are not applied during the sampling. Both quantitative and qualitative experiments are reported.

Strengths

- Motivation is agreeable. - Good writing. - Reasonable method for motivation. - Experiments are done well.

Weaknesses

Specifics of the weaknesses of this paper are written below as questions and limitations, but I believe most of them can be resolved during the rebuttal. I will increase my rating if my concerns can be resolved.

Questions

* The meaning of "FNN" is not defined. * What is the meaning of “(0)” of the delta function in line 138? * Why not let $FNN_{\mu}$ and $FNN_{\sigma}$ take target index as additional input and skip Eq. 7? Once the target index is specified, the matching algorithm may not be needed anymore. * What is the relationship between $q$ in Eq. 7 and stride $s$ and $K$? * It would be better if performance of directly estimating the neighbors is reported as well since it would be more accurate and simple setup to implement the motivation of this paper (even though it has the memory inefficiency, as mentioned in L132). * Can it be described more specifically why KL or JSD cannot be applied (L155)? It is not straightforward to me why KL cannot be applied to this task. * Are both terms in Eq. 3 used together for finetuning? * What happens if Diffusion Models are trained from scratch with the proposed objectives?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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

- The most important part of the paper would be “context prediction” term in Eq . 3. It’s motivation is understandable, but it is not interpreted in terms of optimizing variational bound of negative log likelihood. Is it just “additional” term? or can it be interpreted as a term playing a certain role in maximizing ELBO or likelihood? - This paper is proposing to 1. estimate the neighborhood distribution (instead of directly estimating the neighborhood) and 2. minimize Wasstertein distance as a core objective. Though the authors made an attempt to justify the design choice, I believe it could be compared as a sort of ablation study, which might strengthen the proposed method. --- I found Fig.6 and the first concern is resolved. - I believe training time needs to be compared together in Fig. 7. - Although performance improvement is shown by quantitaive experiments, it is not straightforward how the "context prediction" makes model performance better (qualitatively).

Official Comment by Reviewer DpwQ

Thank authors for their response. My major concerns are well addressed, but I still get confused by some notations in the paper, such as the $h_{t-1}$ and $\delta_{h_u^{(0)}}$, as I mentioned in the review. * How can $h_{t-1}$ represent a point in $x_{t-1}$? By the definition, it is $t-1$-th feature point of the feature map. * Is $\delta$ dirac delta function? Why do you need to denote the center point of $N_i^s$ (i.e., $h_i$) as $h_{u}^{(0)}$? I suggest the authors carefully proofread their manuscript, especially the notations. I will raise my score by 1.