Roadmap

[murrellb]

• [x] Masking - current top priority.
• [ ] Add docs.
• [ ] Tests.
• [x] "Standard" loss functions. One per process (#37).
• [x] Update examples with new interface.
• [x] Add MNIST example (#6).
• [x] Add intractable angular diffusion, and example.
• [x] Remove legacy code remaining from before refactor (#19).
• [x] Process name consistency.
• [x] Add mechanism for allowing "self-conditioning" in the inference loop.
• [x] Generic forwardsample handling of different T per batch.
• [x] Tidy up type system.

[bicycle1885]

I made a PR #6 for the MNIST example.

[murrellb]

"Self-conditioning" feeds the predictions from the previous timestep to the model at the current timestep (https://arxiv.org/pdf/2208.04202.pdf). This can sometimes make a big practical difference. My intuitive understanding is that this gives the model a snapshot of where all the variables are heading to, and helps the model handle conditional dependencies between variables. This needs a modification during training (see eg. https://github.com/lucidrains/denoising-diffusion-pytorch/issues/94), which can be handled by the user, but it also needs a slightly different flow during the reverse diffusion, which we'll need to implement.

[bicycle1885]

To achieve the self-conditioned sampling, we don't need to modify our code if we use a closure trick that looks like this:

function selfconditioned(x)
    x_0 = zero(x)
    function (x_t, t)
        x_0 = net(x_t, x_0, t)
        return x_0
    end
end

x = randn(10)
samplebackward(selfconditioned(x), process, timesteps, x)

[murrellb]

We need to name our processes consistently. Since this is Diffusions.jl, I suggest: OrnsteinUhlenbeckDiffusion RotationDiffusion WrappedBrownianDiffusion UniformDiscreteDiffusion IndependentDiscreteDiffusion etc.