Questions about the third stage

[NwTbbetter]

Hi, I'm having trouble understanding stage 3, can you give a brief description of how stage 3 works and what does continuous_sqrt_alpha_cumprod do once it's passed into the network?

[tiangexiang]

Thanks for your interest in our work! Stage3 is the generation process that tries to 'hallucinate' reasonable and coherent anatomical MRI details on the given MRI acquisition. After determining the location of a MRI slice within the diffusion Markov chain, one can use the Stage1 trained noise model to 'hack' the reverse diffusion process by using input noisy image as an intermediate diffusion posterior. In this way, a clean MRI slice can be inferred by starting from the intermediate state (noisy input slice) only, no need to start from the beginning (pure noise).
We follow a common DDPM implementation to train the diffusion model. continuous_sqrt_alpha_cumprod is used as an indicator of time stamp (which state it is within the Markov chain). As suggested in those common implementations, the denoising U-Net requires both noisy image and a time stamp indicator as inputs to be functional.