About the image reconstruction supervision

[zero0kiriyu]

I am curious about the image reconstruction task you use for the detail preserver. It seems the input of u-net is the concat of noise depth map and rgb token, I wonder how do you prevent the model directly copy the rgb token from input? Can you provide more detail about this supervision?

[haodong2000]

Hi @zero0kiriyu ! Thanks for interets!

Here is the story:

1. The original diffusion model excels at generating detailed images. However, when fine-tuned for dense prediction, it can lose such ability.
2. Fine-tuning to predict both dense labels as well as RGB reconstruction (the pre-training task of diffusion), can mitigate catastrophic forgetting and deliver more precise predictions.

This regularization term is implemented via the interface term "class_labels" of UNet models. It starts with a 2×1 one-hot vector (that acts as a "switcher"), then is the 4×1 vector after positional encoding, the 4×1 vector is then mapped into the same size as the time embedding and is concatenated or added with the time embedding to form the new time embedding. If using concatenation, the dimension of time embedding will be doubled.

[Boatsure]

Sorry, I still don't get it.

Is this 2×1 one-hot vector fixed or a set of learnable parameters? How does this “switcher” control the model's behavior choice of reconstruction or prediction task in training, after all of these position encoding and mapping to the same size as the time-embedding calculation?

In other words, is there a binary judgment logic in the training workflow? Like "if switcher=sx; elif switcher=sy"...

And the reason why the model doesn't directly copy the RGB token from input is that the decoder just cann’t because there is a U-Net in the middle?