One Step Out of the Static

In the lessons, each reverse step takes the noisy input, subtracts most of the predicted noise, and moves one step back along the chain. The original diffusion sampler is a different, random version of that walk. It does one extra thing at every step except the last: after removing noise, it mixes a small amount of fresh random noise back in. Randomness then enters all the way along the walk, not only in the starting static.

The schedule is a list of noise amounts $\beta_1, \beta_2, \dots, \beta_T$, one per step. From it, define

You have the noisy pixels $x_t$ at step $t$, the network's predicted noise $\hat\epsilon$, and a fresh noise draw $z$ (one standard-normal number per pixel). One step back is, for every pixel,

At the very last step, $t = 1$, the fresh-noise term is left out. That step produces the finished image, and nothing is added on top of it.

Task: write reverse_step(x_t, eps_pred, z, betas, t).

• x_t, eps_pred and z are equal-length lists with one entry per pixel.
• betas is the whole schedule: betas[0] is $\beta_1$, betas[1] is $\beta_2$, and so on. It can be longer than t.
• t is the current step, counting from 1.

Return the list $x_{t-1}$, with each value rounded to 4 decimal places.

In a real sampler, z comes from a random-number generator at every step. Here it is passed in so the result can be checked.