How to train score network for mnist?

[youwyu]

Great work. How can we train / get the ann_2999.npz as required by design mnist experiment?

[jcopo]

Thanks! I am currently updating the code of this repo to use https://github.com/jcopo/diffuse. A mnist diffusion model is available here: https://huggingface.co/jcopo/mnist

I will update this repo asap.

[youwyu]

Thank you! Actually I have tried running diffuse/mnist_train code without any issues. The result of ConDiff in design_mnist does not outperform the random policy. I asked LLM to give me a guess of stratified code, could you please tell me whether it's right, since in your referred github repo I also did not find code named with stratified?

def stratified(rng_key, weights, n_particles):
    # Make sure weights are normalized
    weights = weights / jnp.sum(weights)

    # Stratified uniform samples
    positions = (jnp.arange(n_particles) + jax.random.uniform(rng_key, (n_particles,))) / n_particles

    # Cumulative sum of weights defines bins
    cumsum_w = jnp.cumsum(weights)

    # Assign each position to a bin (particle index)
    idx = jnp.searchsorted(cumsum_w, positions, side="right")

    return idx

The first image is from random policy at iteration 10.

The second image is from ConDiff at iteration 10.

In my perspective, the random policy performs slightly better? But I should have mistaken some steps in running your code.

[jcopo]

It seems like the conditioning diffusion is not really working as the missing regions are not properly reconstructed . The conditioning is quite sensitive to how well the diffusion model was trained and the hyperparameters choosen for the guidance.
The optimization is heavily reliant on this conditioning so it doesn't surprise me that it is not working well.

If the conditioning works well, you should have some generated samples similar to this tutorial.

I should be able to update this repo to use diffuse this week. Sorry for the delay!

[youwyu]

Thank you for the reply! No pressure on updating the code. For the diffusion, I have tried the training epochs of 3,500 and 700,000 (clearly it's over-fitting). However, both gives me the similar results. Should I use your diffuse package instead of the diffuse in this repo?