Roadmap

[murrellb]

• Masking - current top priority.
• Add docs.
• Tests.
• "Standard" loss functions. One per process (Standard losses, loss balancing example, rotation example, updating… #37).
• Update examples with new interface.
• Add MNIST example (add MNIST example #6).
• Add intractable angular diffusion, and example.
• Remove legacy code remaining from before refactor (delete legacy code #19).
• Process name consistency.
• Add mechanism for allowing "self-conditioning" in the inference loop.
• Generic forwardsample handling of different T per batch.
• 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. lucidrains/denoising-diffusion-pytorch#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.